Coverage for src/extratools_core/trie.py: 0%
69 statements
« prev ^ index » next coverage.py v7.8.1, created at 2025-05-21 05:55 -0700
« prev ^ index » next coverage.py v7.8.1, created at 2025-05-21 05:55 -0700
1from __future__ import annotations
3from collections.abc import Callable, Iterator, MutableMapping
4from typing import Any
7class TrieDict[VT: Any](MutableMapping):
8 def __init__(self) -> None:
9 self.root: dict[str, Any] = {}
11 self.__len: int = 0
13 def __len__(self) -> int:
14 return self.__len
16 def __find(self, s: str, func: Callable[[dict[str, Any], str], Any]) -> Any:
17 node: dict[str, Any] = self.root
19 while True:
20 c: str = s[0] if s else ""
21 rest: str = s[1:] if s else ""
23 next_node: dict[str, Any] | tuple[str, VT] | None = node.get(c)
24 if next_node is None:
25 raise KeyError
27 if isinstance(next_node, dict):
28 node = next_node
29 s = rest
30 continue
32 if rest == next_node[0]:
33 return func(node, c)
35 raise KeyError
37 def __delitem__(self, s: str) -> None:
38 def delitem(node: dict[str, Any], c: str) -> None:
39 del node[c]
40 self.__len -= 1
42 return self.__find(s, delitem)
44 def __getitem__(self, s: str) -> VT:
45 def getitem(node: dict[str, Any], c: str) -> VT:
46 return node[c][1]
48 return self.__find(s, getitem)
50 def __setitem__(self, s: str, v: VT) -> None:
51 self.__set(s, v, self.root, is_new=True)
53 def __set(self, s: str, v: VT, node: dict[str, Any], *, is_new: bool) -> None:
54 if not s:
55 is_new = is_new and "" not in node
56 node[""] = ("", v)
57 if is_new:
58 self.__len += 1
60 return
62 c: str = s[0]
63 rest: str = s[1:]
65 next_node: dict[str, Any] | tuple[str, VT] | None = node.get(c)
66 if next_node is None:
67 node[c] = (rest, v)
68 if is_new:
69 self.__len += 1
70 elif isinstance(next_node, dict):
71 self.__set(rest, v, next_node, is_new=is_new)
72 else:
73 other_rest: str
74 other_value: VT
75 other_rest, other_value = next_node
77 if rest == other_rest:
78 node[c] = (rest, v)
79 return
81 next_node = node[c] = {}
83 self.__set(other_rest, other_value, next_node, is_new=False)
84 self.__set(rest, v, next_node, is_new=is_new)
86 def __iter__(self) -> Iterator[str]:
87 for _, value in self.__prefixes("", self.root):
88 yield value
90 def prefixes(self) -> Iterator[tuple[str, str]]:
91 yield from self.__prefixes("", self.root)
93 def __prefixes(self, prefix: str, node: dict[str, Any]) -> Iterator[tuple[str, str]]:
94 for key, next_node in node.items():
95 new_prefix = prefix + key
96 if isinstance(next_node, dict):
97 yield from self.__prefixes(new_prefix, next_node)
98 else:
99 yield (new_prefix, new_prefix + next_node[0])