Updates for PR Feedback

Author: strawgate

key-value/key-value-aio/src/key_value/aio/adapters/pydantic/adapter.py

@@ -95,7 +95,13 @@ async def get(self, key: str, *, collection: str | None = None, default: T | Non
 
         return default
 
-    async def get_many(self, keys: Sequence[str], *, collection: str | None = None, default: T | None = None) -> list[T | None]:
+    @overload
+    async def get_many(self, keys: Sequence[str], *, collection: str | None = None, default: T) -> list[T]: ...
+
+    @overload
+    async def get_many(self, keys: Sequence[str], *, collection: str | None = None, default: None = None) -> list[T | None]: ...
+
+    async def get_many(self, keys: Sequence[str], *, collection: str | None = None, default: T | None = None) -> list[T] | list[T | None]:
         """Batch get and validate models by keys, preserving order.
 
         Args:

key-value/key-value-aio/tests/adapters/test_pydantic.py

@@ -38,11 +38,13 @@ class Order(BaseModel):
 FIXED_UPDATED_AT: datetime = datetime(year=2021, month=1, day=1, hour=15, minute=0, second=0, tzinfo=timezone.utc)
 
 SAMPLE_USER: User = User(name="John Doe", email="john.doe@example.com", age=30)
+SAMPLE_USER_2: User = User(name="Jane Doe", email="jane.doe@example.com", age=25)
 SAMPLE_PRODUCT: Product = Product(name="Widget", price=29.99, quantity=10, url=AnyHttpUrl(url="https://example.com"))
 SAMPLE_ORDER: Order = Order(created_at=datetime.now(), updated_at=datetime.now(), user=SAMPLE_USER, product=SAMPLE_PRODUCT, paid=False)
 
 TEST_COLLECTION: str = "test_collection"
 TEST_KEY: str = "test_key"
+TEST_KEY_2: str = "test_key_2"
 
 
 class TestPydanticAdapter:
@@ -80,6 +82,15 @@ async def test_simple_adapter(self, user_adapter: PydanticAdapter[User]):
         cached_user = await user_adapter.get(collection=TEST_COLLECTION, key=TEST_KEY)
         assert cached_user is None
 
+    async def test_simple_adapter_with_default(self, user_adapter: PydanticAdapter[User]):
+        cached_user = await user_adapter.get(collection=TEST_COLLECTION, key=TEST_KEY, default=SAMPLE_USER)
+        assert cached_user == SAMPLE_USER
+
+        await user_adapter.put(collection=TEST_COLLECTION, key=TEST_KEY, value=SAMPLE_USER_2)
+
+        cached_users = await user_adapter.get_many(collection=TEST_COLLECTION, keys=[TEST_KEY, TEST_KEY_2], default=SAMPLE_USER)
+        assert sorted(cached_users, key=lambda x: x.name) == [SAMPLE_USER_2, SAMPLE_USER]
+
     async def test_simple_adapter_with_validation_error_ignore(
         self, user_adapter: PydanticAdapter[User], updated_user_adapter: PydanticAdapter[UpdatedUser]
     ):

key-value/key-value-sync/src/key_value/sync/code_gen/adapters/pydantic/adapter.py

@@ -2,7 +2,7 @@
 # from the original file 'adapter.py'
 # DO NOT CHANGE! Change the original file instead.
 from collections.abc import Sequence
-from typing import Any, Generic, SupportsFloat, TypeVar, get_origin
+from typing import Any, Generic, SupportsFloat, TypeVar, get_origin, overload
 
 from key_value.shared.errors import DeserializationError, SerializationError
 from key_value.shared.type_checking.bear_spray import bear_spray
@@ -65,11 +65,22 @@ def _serialize_model(self, value: T) -> dict[str, Any]:
             msg = f"Invalid Pydantic model: {e}"
             raise SerializationError(msg) from e
 
-    def get(self, key: str, *, collection: str | None = None) -> T | None:
+    @overload
+    def get(self, key: str, *, collection: str | None = None, default: T) -> T: ...
+
+    @overload
+    def get(self, key: str, *, collection: str | None = None, default: None = None) -> T | None: ...
+
+    def get(self, key: str, *, collection: str | None = None, default: T | None = None) -> T | None:
         """Get and validate a model by key.
 
+        Args:
+            key: The key to retrieve.
+            collection: The collection to use. If not provided, uses the default collection.
+            default: The default value to return if the key doesn't exist or validation fails.
+
         Returns:
-            The parsed model instance, or None if not present.
+            The parsed model instance if found and valid, or the default value if key doesn't exist or validation fails.
 
         Raises:
             DeserializationError if the stored data cannot be validated as the model and the PydanticAdapter is configured to
@@ -78,15 +89,28 @@ def get(self, key: str, *, collection: str | None = None) -> T | None:
         collection = collection or self._default_collection
 
         if value := self._key_value.get(key=key, collection=collection):
-            return self._validate_model(value=value)
+            validated = self._validate_model(value=value)
+            if validated is not None:
+                return validated
 
-        return None
+        return default
 
-    def get_many(self, keys: Sequence[str], *, collection: str | None = None) -> list[T | None]:
+    @overload
+    def get_many(self, keys: Sequence[str], *, collection: str | None = None, default: T) -> list[T]: ...
+
+    @overload
+    def get_many(self, keys: Sequence[str], *, collection: str | None = None, default: None = None) -> list[T | None]: ...
+
+    def get_many(self, keys: Sequence[str], *, collection: str | None = None, default: T | None = None) -> list[T] | list[T | None]:
         """Batch get and validate models by keys, preserving order.
 
+        Args:
+            keys: The list of keys to retrieve.
+            collection: The collection to use. If not provided, uses the default collection.
+            default: The default value to return for keys that don't exist or fail validation.
+
         Returns:
-            A list of parsed model instances, or None if missing.
+            A list of parsed model instances, with default values for missing keys or validation failures.
 
         Raises:
             DeserializationError if the stored data cannot be validated as the model and the PydanticAdapter is configured to
@@ -96,7 +120,14 @@ def get_many(self, keys: Sequence[str], *, collection: str | None = None) -> lis
 
         values: list[dict[str, Any] | None] = self._key_value.get_many(keys=keys, collection=collection)
 
-        return [self._validate_model(value=value) if value else None for value in values]
+        result: list[T | None] = []
+        for value in values:
+            if value is None:
+                result.append(default)
+            else:
+                validated = self._validate_model(value=value)
+                result.append(validated if validated is not None else default)
+        return result
 
     def put(self, key: str, value: T, *, collection: str | None = None, ttl: SupportsFloat | None = None) -> None:
         """Serialize and store a model.

key-value/key-value-sync/src/key_value/sync/code_gen/stores/memory/store.py

@@ -2,6 +2,7 @@
 # from the original file 'store.py'
 # DO NOT CHANGE! Change the original file instead.
 import sys
+from collections.abc import Mapping
 from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Any, SupportsFloat
@@ -103,26 +104,63 @@ class MemoryStore(BaseDestroyStore, BaseDestroyCollectionStore, BaseEnumerateCol
     max_entries_per_collection: int
 
     _cache: dict[str, MemoryCollection]
-
-    def __init__(self, *, max_entries_per_collection: int = DEFAULT_MAX_ENTRIES_PER_COLLECTION, default_collection: str | None = None):
+    _seed: Mapping[str, Mapping[str, Mapping[str, Any]]]
+
+    def __init__(
+        self,
+        *,
+        max_entries_per_collection: int = DEFAULT_MAX_ENTRIES_PER_COLLECTION,
+        default_collection: str | None = None,
+        seed: Mapping[str, Mapping[str, Mapping[str, Any]]] | None = None,
+    ):
         """Initialize a fixed-size in-memory store.
 
         Args:
             max_entries_per_collection: The maximum number of entries per collection. Defaults to 10000.
             default_collection: The default collection to use if no collection is provided.
+            seed: Optional seed data to pre-populate the store. Format: {collection: {key: {field: value, ...}}}.
+                Each value must be a mapping (dict) that will be stored as the entry's value.
+                Seeding occurs lazily when each collection is first accessed.
         """
 
         self.max_entries_per_collection = max_entries_per_collection
 
         self._cache = {}
+        self._seed = seed or {}
 
         self._stable_api = True
 
         super().__init__(default_collection=default_collection)
 
+    def _create_collection(self, collection: str) -> MemoryCollection:
+        """Create a new collection.
+
+        Args:
+            collection: The collection name.
+
+        Returns:
+            The created MemoryCollection instance.
+        """
+        collection_cache = MemoryCollection(max_entries=self.max_entries_per_collection)
+        self._cache[collection] = collection_cache
+        return collection_cache
+
     @override
     def _setup_collection(self, *, collection: str) -> None:
-        self._cache[collection] = MemoryCollection(max_entries=self.max_entries_per_collection)
+        """Set up a collection, creating it and seeding it if seed data is available.
+
+        Args:
+            collection: The collection name.
+        """
+        # Create the collection
+        collection_cache = self._create_collection(collection)
+
+        # Seed the collection if seed data is available for it
+        if collection in self._seed:
+            items = self._seed[collection]
+            for key, value in items.items():
+                managed_entry = ManagedEntry(value=value)
+                collection_cache.put(key=key, value=managed_entry)
 
     @override
     def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry | None:

key-value/key-value-sync/src/key_value/sync/code_gen/wrappers/default_value/__init__.py

@@ -0,0 +1,8 @@
+# WARNING: this file is auto-generated by 'build_sync_library.py'
+# from the original file '__init__.py'
+# DO NOT CHANGE! Change the original file instead.
+"""Default value wrapper for returning fallback values when keys are not found."""
+
+from key_value.sync.code_gen.wrappers.default_value.wrapper import DefaultValueWrapper
+
+__all__ = ["DefaultValueWrapper"]

key-value/key-value-sync/src/key_value/sync/code_gen/wrappers/default_value/wrapper.py

@@ -0,0 +1,94 @@
+# WARNING: this file is auto-generated by 'build_sync_library.py'
+# from the original file 'wrapper.py'
+# DO NOT CHANGE! Change the original file instead.
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from typing_extensions import override
+
+from key_value.sync.code_gen.protocols.key_value import KeyValue
+from key_value.sync.code_gen.wrappers.base import BaseWrapper
+
+
+class DefaultValueWrapper(BaseWrapper):
+    """A wrapper that returns a default value when a key is not found.
+
+    This wrapper provides dict.get(key, default) behavior for the key-value store,
+    allowing you to specify a default value to return instead of None when a key doesn't exist.
+    """
+
+    _key_value: KeyValue
+    key_value: KeyValue  # Alias for BaseWrapper compatibility
+
+    def __init__(self, key_value: KeyValue, default_value: Mapping[str, Any], default_ttl: float | None = None) -> None:
+        """Initialize the DefaultValueWrapper.
+
+        Args:
+            key_value: The underlying key-value store to wrap.
+            default_value: The default value to return when a key is not found.
+            default_ttl: The TTL to return for default values. Defaults to None.
+        """
+        self._key_value = key_value
+        self.key_value = key_value  # Alias for BaseWrapper compatibility
+        self._default_value = default_value
+        self._default_ttl = default_ttl
+
+    @override
+    def get(self, key: str, *, collection: str | None = None) -> dict[str, Any] | None:
+        """Get a value by key, returning the default value if not found.
+
+        Args:
+            key: The key to retrieve.
+            collection: The collection to use.
+
+        Returns:
+            The value associated with the key, or the default value if not found.
+        """
+        result = self._key_value.get(key=key, collection=collection)
+        return result if result is not None else dict(self._default_value)
+
+    @override
+    def get_many(self, keys: Sequence[str], *, collection: str | None = None) -> list[dict[str, Any] | None]:
+        """Get multiple values by keys, returning the default value for missing keys.
+
+        Args:
+            keys: The keys to retrieve.
+            collection: The collection to use.
+
+        Returns:
+            A list of values, with default values for missing keys.
+        """
+        results = self._key_value.get_many(keys=keys, collection=collection)
+        return [result if result is not None else dict(self._default_value) for result in results]
+
+    @override
+    def ttl(self, key: str, *, collection: str | None = None) -> tuple[dict[str, Any] | None, float | None]:
+        """Get a value and its TTL, returning the default value if not found.
+
+        Args:
+            key: The key to retrieve.
+            collection: The collection to use.
+
+        Returns:
+            A tuple of (value, ttl), with default value and default TTL if not found.
+        """
+        (result, ttl_value) = self._key_value.ttl(key=key, collection=collection)
+        if result is None:
+            return (dict(self._default_value), self._default_ttl)
+        return (result, ttl_value)
+
+    @override
+    def ttl_many(self, keys: Sequence[str], *, collection: str | None = None) -> list[tuple[dict[str, Any] | None, float | None]]:
+        """Get multiple values with TTLs, returning the default value for missing keys.
+
+        Args:
+            keys: The keys to retrieve.
+            collection: The collection to use.
+
+        Returns:
+            A list of (value, ttl) tuples, with default values and default TTL for missing keys.
+        """
+        results = self._key_value.ttl_many(keys=keys, collection=collection)
+        return [
+            (result, ttl_value) if result is not None else (dict(self._default_value), self._default_ttl) for (result, ttl_value) in results
+        ]

key-value/key-value-sync/src/key_value/sync/code_gen/wrappers/pydantic_json/__init__.py

@@ -0,0 +1,8 @@
+# WARNING: this file is auto-generated by 'build_sync_library.py'
+# from the original file '__init__.py'
+# DO NOT CHANGE! Change the original file instead.
+"""Pydantic JSON wrapper for ensuring JSON-serializable value storage."""
+
+from key_value.sync.code_gen.wrappers.pydantic_json.wrapper import PydanticJsonWrapper
+
+__all__ = ["PydanticJsonWrapper"]

key-value/key-value-sync/src/key_value/sync/code_gen/wrappers/pydantic_json/wrapper.py

@@ -0,0 +1,71 @@
+# WARNING: this file is auto-generated by 'build_sync_library.py'
+# from the original file 'wrapper.py'
+# DO NOT CHANGE! Change the original file instead.
+from collections.abc import Mapping, Sequence
+from typing import Any, SupportsFloat
+
+from pydantic import TypeAdapter
+from typing_extensions import override
+
+from key_value.sync.code_gen.protocols.key_value import KeyValue
+from key_value.sync.code_gen.wrappers.base import BaseWrapper
+
+
+class PydanticJsonWrapper(BaseWrapper):
+    """A wrapper that ensures all values are JSON-serializable using Pydantic's TypeAdapter.
+
+    This wrapper automatically converts values to JSON-safe formats before storage,
+    ensuring compatibility with stores that require JSON-serializable data.
+    """
+
+    _key_value: KeyValue
+    key_value: KeyValue  # Alias for BaseWrapper compatibility
+
+    def __init__(self, key_value: KeyValue) -> None:
+        """Initialize the PydanticJsonWrapper.
+
+        Args:
+            key_value: The underlying key-value store to wrap.
+        """
+        self._key_value = key_value
+        self.key_value = key_value  # Alias for BaseWrapper compatibility
+        self._adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any])
+
+    def _to_json_safe(self, value: Mapping[str, Any]) -> dict[str, Any]:
+        """Convert a value to a JSON-safe format using Pydantic.
+
+        Args:
+            value: The value to convert.
+
+        Returns:
+            A JSON-safe dictionary.
+        """
+        return self._adapter.dump_python(value, mode="json")  # type: ignore[return-value]
+
+    @override
+    def put(self, key: str, value: Mapping[str, Any], *, collection: str | None = None, ttl: SupportsFloat | None = None) -> None:
+        """Store a value after converting it to JSON-safe format.
+
+        Args:
+            key: The key to store.
+            value: The value to store (will be converted to JSON-safe format).
+            collection: The collection to use.
+            ttl: The time-to-live in seconds.
+        """
+        json_safe_value = self._to_json_safe(value)
+        self._key_value.put(key=key, value=json_safe_value, collection=collection, ttl=ttl)
+
+    @override
+    def put_many(
+        self, keys: Sequence[str], values: Sequence[Mapping[str, Any]], *, collection: str | None = None, ttl: SupportsFloat | None = None
+    ) -> None:
+        """Store multiple values after converting them to JSON-safe format.
+
+        Args:
+            keys: The keys to store.
+            values: The values to store (will be converted to JSON-safe format).
+            collection: The collection to use.
+            ttl: The time-to-live in seconds for all items.
+        """
+        json_safe_values = [self._to_json_safe(value) for value in values]
+        self._key_value.put_many(keys=keys, values=json_safe_values, collection=collection, ttl=ttl)

key-value/key-value-sync/src/key_value/sync/wrappers/default_value/__init__.py

@@ -0,0 +1,8 @@
+# WARNING: this file is auto-generated by 'build_sync_library.py'
+# from the original file '__init__.py'
+# DO NOT CHANGE! Change the original file instead.
+"""Default value wrapper for returning fallback values when keys are not found."""
+
+from key_value.sync.code_gen.wrappers.default_value.wrapper import DefaultValueWrapper
+
+__all__ = ["DefaultValueWrapper"]