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

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

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

3

4 from copy import deepcopy

5 from textwrap import indent

6 from typing import Any, Optional, TYPE_CHECKING

7

8 from ..nodes import Node

9 from ..protocols import TrieKeyToken

10 from ..types import TrieEntry, TrieId

11

12 if TYPE_CHECKING:

13 from .trie_mixins import TrieMixinsInterface

14

15

16 class 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: 67 ↛ 72line 67 didn't jump to line 72 because the condition on line 67 was always true

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