Coverage for mcpgateway / plugins / framework / memory.py: 100%
90 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-02-11 07:10 +0000
1# -*- coding: utf-8 -*-
2"""Location: ./mcpgateway/plugins/framework/memory.py
3Copyright 2025
4SPDX-License-Identifier: Apache-2.0
5Authors: Fred Araujo
7Memory management utilities for plugin framework.
9This module provides copy-on-write data structures for efficient memory management
10in plugin contexts.
11"""
13# Standard
14from typing import Any, Iterator, Optional, TypeVar
16T = TypeVar("T")
19class CopyOnWriteDict(dict):
20 """
21 A dictionary subclass that implements copy-on-write behavior.
23 Inherits from dict and layers modifications over an original dictionary
24 without mutating the original. The dict itself stores modifications, while
25 reads check the modifications first, then fall back to the original.
27 This is useful for plugin contexts where you want to isolate modifications
28 without copying the entire original dictionary upfront. Since it subclasses
29 dict, it's compatible with type checking and validation frameworks like Pydantic.
31 Example:
32 >>> original = {"a": 1, "b": 2, "c": 3}
33 >>> cow = CopyOnWriteDict(original)
34 >>> isinstance(cow, dict)
35 True
36 >>> cow["a"] = 10 # Modification stored in dict
37 >>> cow["d"] = 4 # New key stored in dict
38 >>> del cow["b"] # Deletion tracked separately
39 >>> cow["a"]
40 10
41 >>> "b" in cow
42 False
43 >>> original # Original unchanged
44 {'a': 1, 'b': 2, 'c': 3}
45 >>> cow.get_modifications()
46 {'a': 10, 'd': 4}
47 """
49 def __init__(self, original: dict):
50 """
51 Initialize a copy-on-write dictionary wrapper.
53 Args:
54 original: The original dictionary to wrap. This will not be modified.
55 """
56 # Initialize parent dict without any data
57 # The parent dict (self via super()) will store modifications only
58 super().__init__()
59 self._original = original
60 self._deleted = set() # Track keys that have been deleted
62 def __getitem__(self, key: Any) -> Any:
63 """
64 Get an item from the dictionary.
66 Args:
67 key: The key to look up.
69 Returns:
70 The value associated with the key.
72 Raises:
73 KeyError: If the key is not found or has been deleted.
74 """
75 if key in self._deleted:
76 raise KeyError(key)
77 # Check modifications first (via super()), then original
78 if super().__contains__(key):
79 return super().__getitem__(key)
80 if key in self._original:
81 return self._original[key]
82 raise KeyError(key)
84 def __setitem__(self, key: Any, value: Any) -> None:
85 """
86 Set an item in the dictionary.
88 The modification is stored in the wrapper layer, not the original dict.
90 Args:
91 key: The key to set.
92 value: The value to associate with the key.
93 """
94 super().__setitem__(key, value) # Store in modifications (parent dict)
95 self._deleted.discard(key) # If we're setting it, it's not deleted
97 def __delitem__(self, key: Any) -> None:
98 """
99 Delete an item from the dictionary.
101 The key is marked as deleted in the wrapper layer.
103 Args:
104 key: The key to delete.
106 Raises:
107 KeyError: If the key doesn't exist in the dictionary.
108 """
109 if key not in self:
110 raise KeyError(key)
111 self._deleted.add(key)
112 if super().__contains__(key):
113 super().__delitem__(key) # Remove from modifications if present
115 def __contains__(self, key: Any) -> bool:
116 """
117 Check if a key exists in the dictionary.
119 Args:
120 key: The key to check.
122 Returns:
123 True if the key exists and hasn't been deleted, False otherwise.
124 """
125 if key in self._deleted:
126 return False
127 return super().__contains__(key) or key in self._original
129 def __len__(self) -> int:
130 """
131 Get the number of items in the dictionary.
133 Returns:
134 The count of non-deleted keys.
135 """
136 # Get all keys from both modifications and original, excluding deleted
137 all_keys = set(super().keys()) | set(self._original.keys())
138 return len(all_keys - self._deleted)
140 def __iter__(self) -> Iterator:
141 """
142 Iterate over keys in the dictionary.
144 Yields keys in insertion order: first keys from the original dict (in their
145 original order), then new keys from modifications (in their insertion order).
147 Yields:
148 Keys that haven't been deleted.
149 """
150 # First, yield keys from original (in original order)
151 for key in self._original:
152 if key not in self._deleted:
153 yield key
155 # Then yield new keys from modifications (not in original)
156 for key in super().__iter__():
157 if key not in self._original and key not in self._deleted:
158 yield key
160 def __repr__(self) -> str:
161 """
162 Get a string representation of the dictionary.
164 Returns:
165 A string representation showing the current state.
166 """
167 return f"CopyOnWriteDict({dict(self.items())})"
169 def get(self, key: Any, default: Optional[Any] = None) -> Any:
170 """
171 Get an item with a default fallback.
173 Args:
174 key: The key to look up.
175 default: The value to return if the key is not found.
177 Returns:
178 The value associated with the key, or default if not found/deleted.
179 """
180 try:
181 return self[key]
182 except KeyError:
183 return default
185 def keys(self):
186 """
187 Get all non-deleted keys.
189 Returns:
190 A generator of keys.
191 """
192 return iter(self)
194 def values(self):
195 """
196 Get all values for non-deleted keys.
198 Returns:
199 A generator of values.
200 """
201 return (self[k] for k in self)
203 def items(self):
204 """
205 Get all key-value pairs for non-deleted keys.
207 Returns:
208 A generator of (key, value) tuples.
209 """
210 return ((k, self[k]) for k in self)
212 def copy(self) -> dict:
213 """
214 Create a regular dictionary with all current key-value pairs.
216 Returns:
217 A new dict containing the current state (original + modifications - deletions).
218 """
219 return dict(self.items())
221 def get_modifications(self) -> dict:
222 """
223 Get only the modifications made to the wrapper.
225 This returns only the keys that were added or changed in the modification layer,
226 not including values from the original dictionary that weren't modified.
228 Returns:
229 A copy of the modifications dictionary.
230 """
231 # The parent dict (super()) contains only modifications
232 return dict(super().items())
234 def get_deleted(self) -> set:
235 """
236 Get the set of deleted keys.
238 Returns:
239 A copy of the deleted keys set.
240 """
241 return self._deleted.copy()
243 def has_modifications(self) -> bool:
244 """
245 Check if any modifications have been made.
247 Returns:
248 True if there are any modifications or deletions, False otherwise.
249 """
250 # Check if parent dict has any entries (modifications) or if anything was deleted
251 return super().__len__() > 0 or len(self._deleted) > 0
253 def update(self, other=None, **kwargs) -> None:
254 """
255 Update the dictionary with key-value pairs from another mapping or iterable.
257 Args:
258 other: A mapping or iterable of key-value pairs.
259 **kwargs: Additional key-value pairs to update.
261 Examples:
262 >>> cow = CopyOnWriteDict({"a": 1})
263 >>> cow.update({"b": 2, "c": 3})
264 >>> cow.update(d=4, e=5)
265 >>> dict(cow.items())
266 {'a': 1, 'b': 2, 'c': 3, 'd': 4, 'e': 5}
267 """
268 if other is not None:
269 if hasattr(other, "items"):
270 for key, value in other.items():
271 self[key] = value
272 else:
273 for key, value in other:
274 self[key] = value
275 for key, value in kwargs.items():
276 self[key] = value
278 def pop(self, key: Any, *args) -> Any:
279 """
280 Remove and return the value for a key.
282 Args:
283 key: The key to remove.
284 *args: Optional default value if key is not found.
286 Returns:
287 The value associated with the key.
289 Raises:
290 KeyError: If key is not found and no default is provided.
291 TypeError: If more than one default argument is provided.
293 Examples:
294 >>> cow = CopyOnWriteDict({"a": 1, "b": 2})
295 >>> cow.pop("a")
296 1
297 >>> cow.pop("c", "default")
298 'default'
299 """
300 if len(args) > 1:
301 raise TypeError(f"pop() accepts 1 or 2 arguments ({len(args) + 1} given)")
303 try:
304 value = self[key]
305 del self[key]
306 return value
307 except KeyError:
308 if args:
309 return args[0]
310 raise
312 def setdefault(self, key: Any, default: Any = None) -> Any:
313 """
314 Get a value, setting it to a default if not present.
316 Args:
317 key: The key to look up.
318 default: The default value to set if key is not present.
320 Returns:
321 The value associated with the key (existing or newly set).
323 Examples:
324 >>> cow = CopyOnWriteDict({"a": 1})
325 >>> cow.setdefault("a", 10)
326 1
327 >>> cow.setdefault("b", 2)
328 2
329 >>> cow["b"]
330 2
331 """
332 if key in self:
333 return self[key]
334 self[key] = default
335 return default
337 def clear(self) -> None:
338 """
339 Remove all items from the dictionary.
341 This marks all keys (from original and modifications) as deleted.
343 Examples:
344 >>> cow = CopyOnWriteDict({"a": 1, "b": 2})
345 >>> cow.clear()
346 >>> len(cow)
347 0
348 """
349 # Mark all current keys as deleted
350 for key in list(self.keys()):
351 self._deleted.add(key)
352 # Clear modifications from parent dict
353 super().clear()
356def copyonwrite(o: T) -> T:
357 """
358 Returns a copy-on-write wrapper of the original object.
360 Args:
361 o: The object to wrap. Currently only supports dict objects.
363 Returns:
364 A copy-on-write wrapper around the object.
366 Raises:
367 TypeError: If the object type is not supported for copy-on-write wrapping.
368 """
369 if isinstance(o, dict):
370 return CopyOnWriteDict(o)
371 raise TypeError(f"No copy-on-write wrapper available for {type(o)}")