Coverage for src/gentrie/trie/storage.py: 100%
34 statements
« prev ^ index » next coverage.py v7.6.10, created at 2025-08-17 11:24 -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 # Traverse the trie to find the insertion point for the key,
117 # creating nodes as necessary.
118 current_node = self
119 for token in key:
120 if token not in current_node.children:
121 child_node = Node(token=token, parent=current_node) # type: ignore[reportArgumentType]
122 current_node.children[token] = child_node
123 current_node = current_node.children[token]
125 # This key is already in the trie (it has a trie id)
126 if current_node.ident:
127 # If we allow updating, update the value and return the existing id
128 if allow_value_update:
129 current_node.value = value
130 self._trie_entries[current_node.ident] = TrieEntry(
131 current_node.ident, key, value)
132 return current_node.ident
134 # The key is already in the trie but we are not allowing updating values - so raise an error
135 raise DuplicateKeyError(
136 msg=("Attempted to store a key with a value that is already in the trie "
137 " - use `update()` to change the value of an existing key."),
138 tag=ErrorTag.STORE_ENTRY_DUPLICATE_KEY)
140 # Assign a new trie id for the node and set the value
141 self._ident_counter += 1
142 new_ident = TrieId(self._ident_counter)
143 current_node.ident = new_ident
144 current_node.value = value
145 self._trie_index[new_ident] = cast(Node, current_node)
146 self._trie_entries[new_ident] = TrieEntry(new_ident, key, value)
148 return new_ident