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

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

2"""Data access operations for the trie."""

3

4 from typing import Any, Optional

5

6 from ..exceptions import ErrorTag, InvalidGeneralizedKeyError, TrieKeyError, TrieTypeError

7 from ..types import TrieEntry, TrieId, GeneralizedKey

8 from ..validation import is_generalizedkey

9

10 from .trie_mixins import TrieMixinsInterface

11

12# Disabled because pyright does not understand mixins

13# use of private attributes from the mixing class as declared

14# in the TrieMixinsInterface protocol.

15# pyright: reportPrivateUsage=false

16

17

18 class TrieAccessMixin:

19 """Mixin providing data access operations.

20

21 Note: This mixin accesses private attributes of the mixing class.

22 This is intentional and necessary for the mixin pattern.

23 """

24 def __contains__(self: TrieMixinsInterface, key_or_ident: GeneralizedKey | TrieId) -> bool:

25 """Returns True if the trie contains a GeneralizedKey or TrieId matching the passed key.

26

27 This method checks if the trie contains a key that matches the provided key_or_ident.

28 The key can be specified either as a :class:`GeneralizedKey` or as a :class:`TrieId`.

29

30 A lookup by :class:`TrieId` is a fast operation (*O(1)* time) while a lookup by :class:`GeneralizedKey`

31 involves traversing the trie structure to find a matching key (*O(n)* time in the worst case,

32 where n is the key length).

33

34 Args:

35 key_or_ident (GeneralizedKey | TrieId): Key or TrieId for matching.

36

37 Returns:

38 :class:`bool`: True if there is a matching GeneralizedKey/TrieId in the trie. False otherwise.

39 """

40 if isinstance(key_or_ident, TrieId):

41 # If it's a TrieId, check if it exists in the trie index

42 return key_or_ident in self._trie_index

43

44 if self.runtime_validation and not is_generalizedkey(key_or_ident):

45 return False

46

47 current_node = self

48 for token in key_or_ident:

49 if token not in current_node.children:

50 return False

51 current_node = current_node.children[token]

52

53 return current_node.ident is not None

54

55 def __getitem__(self: TrieMixinsInterface, key: TrieId | GeneralizedKey) -> TrieEntry:

56 """Returns the :class:`TrieEntry` for the ident or key with the passed identifier.

57

58 The identifier can be either the :class:`TrieId` (ident) or the :class:`GeneralizedKey` (key)

59 for the entry.

60

61 Args:

62 key (TrieId | GeneralizedKey): the identifier to retrieve.

63

64 Returns: :class:`TrieEntry`: TrieEntry for the key with the passed identifier.

65

66 Raises:

67 TrieKeyError: if the key arg does not match any keys/idents in the trie.

68 TrieTypeError: if the key arg is neither a :class:`TrieId` or a valid :class:`GeneralizedKey`.

69 """

70 if isinstance(key, TrieId):

71 if key not in self._trie_index:

72 raise TrieKeyError("TrieId not found in trie index", ErrorTag.GETITEM_ID_NOT_FOUND)

73 # Return the TrieEntry for the TrieId

74 return self._trie_entries[key]

75

76 # If runtime validation is disabled, we just ASSUME the key is a GeneralizedKey if we get here.

77 if (not self.runtime_validation) or is_generalizedkey(key):

78 # Find the TrieId for the key

79 current_node = self

80 for token in key:

81 if token not in current_node.children:

82 raise TrieKeyError(

83 msg="key does not match any idents or keys in the trie",

84 tag=ErrorTag.GETITEM_KEY_NOT_FOUND,

85 )

86 current_node = current_node.children[token]

87 if current_node.ident:

88 # Return the TrieEntry for the TrieId

89 return self._trie_entries[current_node.ident]

90 raise TrieKeyError(

91 msg="key does not match any idents or keys in the trie",

92 tag=ErrorTag.GETITEM_NOT_TERMINAL,

93 )

94

95 # If we reach here, the passed key was neither a TrieId nor a GeneralizedKey

96 raise TrieTypeError(

97 msg="key must be either a :class:TrieId or a :class:`GeneralizedKey`",

98 tag=ErrorTag.GETITEM_INVALID_KEY_TYPE

99 )

100

101 def get(self: TrieMixinsInterface,

102 key: TrieId | GeneralizedKey,

103 default: Optional[Any] = None) -> Optional[TrieEntry | Any]:

104 """Returns the :class:`TrieEntry` for the ident or key with the passed identifier.

105

106 The identifier can be either the :class:`TrieId` (ident) or the :class:`GeneralizedKey` (key)

107 for the entry.

108

109 If the key is not found, it returns the default value if provided or None if not provided.

110

111 Args:

112 key (TrieId | GeneralizedKey): the identifier to retrieve.

113 default (Optional[TrieEntry | Any], default=None): The default value to return if the key is not found.

114

115 Returns: :class:`TrieEntry`: TrieEntry for the key with the passed identifier or the default value if not found.

116 """

117 try:

118 return self[key]

119 except (TrieKeyError, TrieTypeError, InvalidGeneralizedKeyError):

120 return default