Coverage for src/gentrie/trie/base.py: 100%

1# -*- coding: utf-8 -*- 

2"""Base trie functionality and initialization.""" 

3 

4from copy import deepcopy 

5from textwrap import indent 

6from typing import Any, Optional, TYPE_CHECKING 

7 

8from ..nodes import Node 

9from ..protocols import TrieKeyToken 

10from ..types import TrieEntry, TrieId 

11 

12if TYPE_CHECKING: 

13 from .trie_mixins import TrieMixinsInterface 

14 

15 

16class TrieBase: 

17 """Base class providing core trie structure and utilities. 

18 

19 Properties: 

20 runtime_validation (bool): Whether to enable runtime validation of keys. 

21 """ 

22 

23 def __init__(self, runtime_validation: bool = True) -> None: 

24 """Initializes a new TrieBase instance. 

25 

26 By default, runtime validation of keys is enabled. If your code is well tested, 

27 you can disable it for improved performance. How the code will react to invalid 

28 keys when validation is disabled is not defined and may lead to unexpected 

29 behavior. 

30 

31 Args: 

32 runtime_validation (bool, default=True): Whether to enable runtime validation of keys. 

33 

34 """ 

35 self.runtime_validation: bool = bool(runtime_validation) 

36 self.token: Optional[TrieKeyToken] = None 

37 self.value: Optional[Any] = None 

38 # The parent of the root node is always None. 

39 # Typing it as the protocol interface makes it compatible with the 

40 # final GeneralizedTrie class. 

41 self.parent: Optional["TrieMixinsInterface"] = None 

42 self.children: dict[TrieKeyToken, Node] = {} 

43 self.ident: Optional[TrieId] = None 

44 # Counter for the next unique identifier to assign to a key in the trie. 

45 self._ident_counter: int = 0 

46 # Mapping of unique identifiers to their corresponding trie nodes. 

47 self._trie_index: dict[TrieId, Node] = {} 

48 # Mapping of unique identifiers to their corresponding TrieEntry instances. 

49 self._trie_entries: dict[TrieId, TrieEntry] = {} 

50 

51 def clear(self) -> None: 

52 """Clears all keys from the trie.""" 

53 self.ident = None 

54 self.token = None 

55 self.value = None 

56 self.parent = None 

57 self.children.clear() 

58 self._trie_index.clear() 

59 self._trie_entries.clear() 

60 # Reset the ident counter 

61 self._ident_counter = 0 

62 

63 def __str__(self) -> str: 

64 """Generates a stringified version of the trie for visual examination.""" 

65 output: list[str] = ["{"] 

66 output.append(f" trie number = {self._ident_counter}") 

67 if self.children: 

68 output.append(" children = {") 

69 for child_key, child_value in self.children.items(): 

70 output.append(f" {repr(child_key)} = " + indent(str(child_value), " ").lstrip()) 

71 output.append(" }") 

72 output.append(f" trie index = {self._trie_index.keys()}") 

73 output.append("}") 

74 return "\n".join(output) 

75 

76 def _as_dict(self) -> dict[str, Any]: 

77 """Converts the trie to a dictionary representation.""" 

78 # pylint: disable=protected-access, no-member 

79 return deepcopy({ 

80 "ident": self.ident, 

81 "children": { 

82 k: v._as_dict() # type: ignore[reportPrivateUsage] 

83 for k, v in self.children.items() 

84 }, 

85 "trie_index": sorted(self._trie_index.keys()), 

86 "trie_entries": { 

87 k: repr(v) 

88 for k, v in self._trie_entries.items() 

89 } 

90 })