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

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

2"""Entry storage operations for the trie."""

3

4# pylint does not understand the Protocol declaration of the mixing class

5# pylint: disable=no-member

6

7 from typing import Any, cast, Optional

8

9 from ..exceptions import DuplicateKeyError, InvalidGeneralizedKeyError, ErrorTag

10 from ..nodes import Node

11 from ..protocols import GeneralizedKey

12 from ..types import TrieEntry, TrieId

13 from ..validation import is_generalizedkey

14

15 from .trie_mixins import TrieMixinsInterface

16

17# Disabled because pyright does not understand mixins

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

19# in the TrieMixinsInterface protocol.

20# pyright: reportPrivateUsage=false

21

22

23 class TrieStorageMixin:

24 """Mixin providing entry storage operations."""

25

26 def add(self: TrieMixinsInterface, key: GeneralizedKey, value: Optional[Any] = None) -> TrieId:

27 """Adds the key to the trie.

28

29 .. warning:: **Keys Must Be Immutable**

30

31 Once a key is added to the trie, neither the key sequence itself nor any of its

32 constituent tokens should be mutated. Modifying a key after it has been added

33 can corrupt the internal state of the trie, leading to unpredictable behavior

34 and making entries unreachable. The trie does not create a deep copy of keys

35 for performance reasons.

36

37 If you need to modify a key, you should remove the old key and add a new one

38 with the modified value.

39

40 Args:

41 key (GeneralizedKey): Must be an object that can be iterated and that when iterated

42 returns elements conforming to the :class:`TrieKeyToken` protocol.

43 value (Optional[Any], default=None): Optional value to associate with the key.

44

45 Raises:

46 InvalidGeneralizedKeyError:

47 If key is not a valid :class:`GeneralizedKey`.

48 DuplicateKeyError:

49 If the key is already in the trie.

50

51 Returns:

52 :class:`TrieId`: Id of the inserted key. If the key was not in the trie,

53 it returns the id of the new entry. If the key was already in the trie,

54 it raises a :class:`DuplicateKeyError`.

55 """

56 return self._store_entry(key=key, value=value, allow_value_update=False)

57

58 def update(self: TrieMixinsInterface, key: GeneralizedKey, value: Optional[Any] = None) -> TrieId:

59 """Updates the key/value pair in the trie.

60

61 .. warning:: **Keys Must Be Immutable**

62

63 Once a key is added to the trie, neither the key sequence itself nor any of its

64 constituent tokens should be mutated. Modifying a key after it has been added

65 can corrupt the internal state of the trie, leading to unpredictable behavior

66 and making entries unreachable. The trie does not create a deep copy of keys

67 for performance reasons.

68

69 If you need to modify a key, you should remove the old key and add a new one

70 with the modified value.

71

72 Args:

73 key (GeneralizedKey): Must be an object that can be iterated and that when iterated

74 returns elements conforming to the :class:`TrieKeyToken` protocol.

75 value (Optional[Any], default=None): Optional value to associate with the key.

76

77 Raises:

78 InvalidGeneralizedKeyError:

79 If key is not a valid :class:`GeneralizedKey`.

80

81 Returns:

82 :class:`TrieId`: Id of the inserted key. If the key was already in the trie with the same value

83 it returns the id for the already existing entry. If the key was not already in the trie,

84 it returns the id for a new entry.

85 """

86 return self._store_entry(key=key, value=value, allow_value_update=True)

87

88 def _store_entry(self: TrieMixinsInterface,

89 key: GeneralizedKey,

90 value: Any,

91 allow_value_update: bool) -> TrieId:

92 """Stores a key/value pair entry in the trie.

93

94 Args:

95 key (GeneralizedKey): Must be an object that can be iterated and that when iterated

96 returns elements conforming to the :class:`TrieKeyToken` protocol.

97 value (Optional[Any], default=None): Optional value to associate with the key.

98 allow_value_update (bool):

99 Whether to allow overwriting the value if the key already exists.

100 Raises:

101 InvalidGeneralizedKeyError: If key is not a valid :class:`GeneralizedKey`.

102 DuplicateKeyError: If the key is already in the trie and `allow_value_update` is False.

103

104 Returns:

105 :class:`TrieId`: Id of the key's entry. If the key was not in the trie,

106 it returns the id of the new entry. If the key was already in the trie and

107 `allow_value_update` is True, it updates the value and returns the existing id.

108 If the key was already in the trie and `allow_value_update` is False,

109 it raises a :class:`DuplicateKeyError`.

110 """

111 if self.runtime_validation and not is_generalizedkey(key):

112 raise InvalidGeneralizedKeyError(

113 msg="key is not a valid `GeneralizedKey`",

114 tag=ErrorTag.STORE_ENTRY_INVALID_GENERALIZED_KEY)

115

116 # convert potentially mutable types of Sequences to immutable types

117 if not isinstance(key, (str, tuple, bytes, frozenset)):

118 key = tuple(key)

119

120 # Traverse the trie to find the insertion point for the key,

121 # creating nodes as necessary.

122 current_node = self

123 for token in key:

124 if token not in current_node.children:

125 child_node = Node(token=token, parent=current_node) # type: ignore[reportArgumentType]

126 current_node.children[token] = child_node

127 current_node = current_node.children[token]

128

129 # This key is already in the trie (it has a trie id)

130 if current_node.ident:

131 # If we allow updating, update the value and return the existing id

132 if allow_value_update:

133 current_node.value = value

134 self._trie_entries[current_node.ident] = TrieEntry(

135 current_node.ident, key, value)

136 return current_node.ident

137

138 # The key is already in the trie but we are not allowing updating values - so raise an error

139 raise DuplicateKeyError(

140 msg=("Attempted to store a key with a value that is already in the trie "

141 " - use `update()` to change the value of an existing key."),

142 tag=ErrorTag.STORE_ENTRY_DUPLICATE_KEY)

143

144 # Assign a new trie id for the node and set the value

145 # If the key is not already an immutable builtin Sequence type, the

146 # key is converted to a tuple to reduce its vulnerability

147 # to changes before being stored.

148 self._ident_counter += 1

149 new_ident = TrieId(self._ident_counter)

150 current_node.ident = new_ident

151 current_node.value = value

152 self._trie_index[new_ident] = cast(Node, current_node)

153 self._trie_entries[new_ident] = TrieEntry(new_ident, key, value)

154

155 return new_ident