Coverage for src/gentrie/trie/storage.py: 100%
36 statements
« prev ^ index » next coverage.py v7.6.10, created at 2025-08-20 09:30 -0700
1# -*- coding: utf-8 -*-
2"""Entry storage operations for the trie."""
4# pylint does not understand the Protocol declaration of the mixing class
5# pylint: disable=no-member
7from typing import Any, cast, Optional
9from ..exceptions import DuplicateKeyError, InvalidGeneralizedKeyError, ErrorTag
10from ..nodes import Node
11from ..protocols import GeneralizedKey
12from ..types import TrieEntry, TrieId
13from ..validation import is_generalizedkey
15from .trie_mixins import TrieMixinsInterface
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
23class TrieStorageMixin:
24 """Mixin providing entry storage operations."""
26 def add(self: TrieMixinsInterface, key: GeneralizedKey, value: Optional[Any] = None) -> TrieId:
27 """Adds the key to the trie.
29 .. warning:: **Keys Must Be Immutable**
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.
37 If you need to modify a key, you should remove the old key and add a new one
38 with the modified value.
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.
45 Raises:
46 InvalidGeneralizedKeyError:
47 If key is not a valid :class:`GeneralizedKey`.
48 DuplicateKeyError:
49 If the key is already in the trie.
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)
58 def update(self: TrieMixinsInterface, key: GeneralizedKey, value: Optional[Any] = None) -> TrieId:
59 """Updates the key/value pair in the trie.
61 .. warning:: **Keys Must Be Immutable**
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.
69 If you need to modify a key, you should remove the old key and add a new one
70 with the modified value.
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.
77 Raises:
78 InvalidGeneralizedKeyError:
79 If key is not a valid :class:`GeneralizedKey`.
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)
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.
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.
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)
116 # convert potentially mutable types of Sequences to immutable types
117 if not isinstance(key, (str, tuple, bytes, frozenset)):
118 key = tuple(key)
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]
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
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)
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)
155 return new_ident