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

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

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

3 

4from typing import Any, Optional 

5 

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

7from ..types import TrieEntry, TrieId, GeneralizedKey 

8from ..validation import is_generalizedkey 

9 

10from .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 

18class 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