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

1from __future__ import annotations 

2 

3from collections.abc import Callable, Iterator, MutableMapping 

4from typing import Any 

5 

6 

7class TrieDict[VT: Any](MutableMapping): 

8 def __init__(self) -> None: 

9 self.root: dict[str, Any] = {} 

10 

11 self.__len: int = 0 

12 

13 def __len__(self) -> int: 

14 return self.__len 

15 

16 def __find(self, s: str, func: Callable[[dict[str, Any], str], Any]) -> Any: 

17 node: dict[str, Any] = self.root 

18 

19 while True: 

20 c: str = s[0] if s else "" 

21 rest: str = s[1:] if s else "" 

22 

23 next_node: dict[str, Any] | tuple[str, VT] | None = node.get(c) 

24 if next_node is None: 

25 raise KeyError 

26 

27 if isinstance(next_node, dict): 

28 node = next_node 

29 s = rest 

30 continue 

31 

32 if rest == next_node[0]: 

33 return func(node, c) 

34 

35 raise KeyError 

36 

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 

41 

42 return self.__find(s, delitem) 

43 

44 def __getitem__(self, s: str) -> VT: 

45 def getitem(node: dict[str, Any], c: str) -> VT: 

46 return node[c][1] 

47 

48 return self.__find(s, getitem) 

49 

50 def __setitem__(self, s: str, v: VT) -> None: 

51 self.__set(s, v, self.root, is_new=True) 

52 

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 

59 

60 return 

61 

62 c: str = s[0] 

63 rest: str = s[1:] 

64 

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 

76 

77 if rest == other_rest: 

78 node[c] = (rest, v) 

79 return 

80 

81 next_node = node[c] = {} 

82 

83 self.__set(other_rest, other_value, next_node, is_new=False) 

84 self.__set(rest, v, next_node, is_new=is_new) 

85 

86 def __iter__(self) -> Iterator[str]: 

87 for _, value in self.__prefixes("", self.root): 

88 yield value 

89 

90 def prefixes(self) -> Iterator[tuple[str, str]]: 

91 yield from self.__prefixes("", self.root) 

92 

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])