Run sync code-gen

Author: strawgate

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

@@ -50,6 +50,21 @@ def __init__(
         self._raise_on_validation_error = raise_on_validation_error
 
     def _validate_model(self, value: dict[str, Any]) -> T | None:
+        """Validate and deserialize a dict into the configured Pydantic model.
+
+        This method handles both single models and list models. For list models, it expects the value
+        to contain an "items" key with the list data, following the convention used by `_serialize_model`.
+        If validation fails and `raise_on_validation_error` is False, returns None instead of raising.
+
+        Args:
+            value: The dict to validate and convert to a Pydantic model.
+
+        Returns:
+            The validated model instance, or None if validation fails and errors are suppressed.
+
+        Raises:
+            DeserializationError: If validation fails and `raise_on_validation_error` is True.
+        """
         try:
             if self._is_list_model:
                 return self._type_adapter.validate_python(value.get("items", []))
@@ -62,6 +77,22 @@ def _validate_model(self, value: dict[str, Any]) -> T | None:
             return None
 
     def _serialize_model(self, value: T) -> dict[str, Any]:
+        """Serialize a Pydantic model to a dict for storage.
+
+        This method handles both single models and list models. For list models, it wraps the serialized
+        list in a dict with an "items" key (e.g., {"items": [...]}) to ensure consistent dict-based storage
+        format across all value types. This wrapping convention is expected by `_validate_model` during
+        deserialization.
+
+        Args:
+            value: The Pydantic model instance to serialize.
+
+        Returns:
+            A dict representation of the model suitable for storage.
+
+        Raises:
+            SerializationError: If the model cannot be serialized.
+        """
         try:
             if self._is_list_model:
                 return {"items": self._type_adapter.dump_python(value, mode="json")}
@@ -91,6 +122,10 @@ def get(self, key: str, *, collection: str | None = None, default: T | None = No
         Raises:
             DeserializationError if the stored data cannot be validated as the model and the PydanticAdapter is configured to
             raise on validation error.
+
+        Note:
+            When raise_on_validation_error=False and validation fails, returns the default value (which may be None).
+            When raise_on_validation_error=True and validation fails, raises DeserializationError.
         """
         collection = collection or self._default_collection
 
@@ -121,6 +156,11 @@ def get_many(self, keys: Sequence[str], *, collection: str | None = None, defaul
         Raises:
             DeserializationError if the stored data cannot be validated as the model and the PydanticAdapter is configured to
             raise on validation error.
+
+        Note:
+            When raise_on_validation_error=False and validation fails for any key, that position in the returned list
+            will contain the default value (which may be None). The method returns a complete list matching the order
+            and length of the input keys, with defaults substituted for missing or invalid entries.
         """
         collection = collection or self._default_collection
 
@@ -171,7 +211,16 @@ def delete_many(self, keys: Sequence[str], *, collection: str | None = None) ->
     def ttl(self, key: str, *, collection: str | None = None) -> tuple[T | None, float | None]:
         """Get a model and its TTL seconds if present.
 
-        Returns (model, ttl_seconds) or (None, None) if missing.
+        Args:
+            key: The key to retrieve.
+            collection: The collection to use. If not provided, uses the default collection.
+
+        Returns:
+            A tuple of (model, ttl_seconds). Returns (None, None) if the key is missing or validation fails.
+
+        Note:
+            When validation fails and raise_on_validation_error=False, returns (None, None) even if TTL data exists.
+            When validation fails and raise_on_validation_error=True, raises DeserializationError.
         """
         collection = collection or self._default_collection
 

key-value/key-value-sync/src/key_value/sync/code_gen/protocols/key_value.py

@@ -7,7 +7,12 @@
 
 @runtime_checkable
 class KeyValueProtocol(Protocol):
-    """A subset of KV operations: get/put/delete and TTL variants, including bulk calls."""
+    """A subset of KV operations: get/put/delete and TTL variants, including bulk calls.
+
+    This protocol defines the minimal contract for key-value store implementations. All methods are
+    async and may raise exceptions on connection failures, validation errors, or other operational issues.
+    Implementations should handle backend-specific errors appropriately.
+    """
 
     def get(self, key: str, *, collection: str | None = None) -> dict[str, Any] | None:
         """Retrieve a value by key from the specified collection.
@@ -52,6 +57,9 @@ def delete(self, key: str, *, collection: str | None = None) -> bool:
         Args:
             key: The key to delete the value from.
             collection: The collection to delete the value from. If no collection is provided, it will use the default collection.
+
+        Returns:
+            True if the key was deleted, False if the key did not exist.
         """
         ...
 

key-value/key-value-sync/src/key_value/sync/code_gen/stores/base.py

@@ -35,6 +35,18 @@
 
 
 def _seed_to_frozen_seed_data(seed: SEED_DATA_TYPE) -> FROZEN_SEED_DATA_TYPE:
+    """Convert mutable seed data to an immutable frozen structure.
+
+    This function converts the nested mapping structure of seed data into immutable MappingProxyType
+    objects at all levels. Using immutable structures prevents accidental modification of seed data
+    after store initialization and ensures thread-safety.
+
+    Args:
+        seed: The mutable seed data mapping: {collection: {key: {field: value}}}.
+
+    Returns:
+        An immutable frozen version of the seed data using MappingProxyType.
+    """
     return MappingProxyType(
         {
             collection: MappingProxyType({key: MappingProxyType(value) for (key, value) in items.items()})
@@ -107,6 +119,15 @@ def _seed_store(self) -> None:
                 self.put(key=key, value=dict(value), collection=collection)
 
     def setup(self) -> None:
+        """Initialize the store if not already initialized.
+
+        This method is called automatically before any store operations and uses a lock to ensure
+        thread-safe lazy initialization. It can also be called manually to ensure the store is ready
+        before performing operations. The setup process includes calling the `_setup()` hook and
+        seeding the store with initial data if provided.
+
+        This method is idempotent - calling it multiple times has no additional effect after the first call.
+        """
         if not self._setup_complete:
             with self._setup_lock:
                 if not self._setup_complete:
@@ -122,6 +143,19 @@ def setup(self) -> None:
                     self._seed_store()
 
     def setup_collection(self, *, collection: str) -> None:
+        """Initialize a specific collection if not already initialized.
+
+        This method is called automatically before any collection-specific operations and uses a per-collection
+        lock to ensure thread-safe lazy initialization. It can also be called manually to ensure a collection
+        is ready before performing operations on it. The setup process includes calling the `_setup_collection()`
+        hook for store-specific collection initialization.
+
+        This method is idempotent - calling it multiple times for the same collection has no additional effect
+        after the first call.
+
+        Args:
+            collection: The name of the collection to initialize.
+        """
         self.setup()
 
         if not self._setup_collection_complete[collection]:

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

@@ -59,6 +59,19 @@
 
 
 def managed_entry_to_document(collection: str, key: str, managed_entry: ManagedEntry) -> dict[str, Any]:
+    """Convert a ManagedEntry to an Elasticsearch document.
+
+    This function creates an Elasticsearch document containing the collection, key,
+    JSON-serialized value, and optional timestamp metadata.
+
+    Args:
+        collection: The collection name.
+        key: The entry key.
+        managed_entry: The ManagedEntry to convert.
+
+    Returns:
+        An Elasticsearch document dictionary.
+    """
     document: dict[str, Any] = {"collection": collection, "key": key, "value": managed_entry.to_json(include_metadata=False)}
 
     if managed_entry.created_at:
@@ -70,6 +83,20 @@ def managed_entry_to_document(collection: str, key: str, managed_entry: ManagedE
 
 
 def source_to_managed_entry(source: dict[str, Any]) -> ManagedEntry:
+    """Convert an Elasticsearch document source to a ManagedEntry.
+
+    This function deserializes an Elasticsearch document back to a ManagedEntry,
+    parsing the JSON-encoded value and timestamp metadata.
+
+    Args:
+        source: The Elasticsearch document _source dictionary.
+
+    Returns:
+        A ManagedEntry reconstructed from the document.
+
+    Raises:
+        DeserializationError: If the value field is missing or not a valid string.
+    """
     if not (value_str := source.get("value")) or not isinstance(value_str, str):
         msg = "Value is not a string"
         raise DeserializationError(msg)

key-value/key-value-sync/src/key_value/sync/code_gen/stores/elasticsearch/utils.py

@@ -8,6 +8,14 @@
 
 
 def get_body_from_response(response: ObjectApiResponse[Any]) -> dict[str, Any]:
+    """Extract and validate the body from an Elasticsearch response.
+
+    Args:
+        response: The Elasticsearch API response object.
+
+    Returns:
+        The response body as a dictionary, or an empty dict if the body is missing or invalid.
+    """
     if not (body := response.body):  # pyright: ignore[reportAny]
         return {}
 
@@ -18,6 +26,14 @@ def get_body_from_response(response: ObjectApiResponse[Any]) -> dict[str, Any]:
 
 
 def get_source_from_body(body: dict[str, Any]) -> dict[str, Any]:
+    """Extract and validate the _source field from an Elasticsearch response body.
+
+    Args:
+        body: The response body dictionary from Elasticsearch.
+
+    Returns:
+        The _source field as a dictionary, or an empty dict if missing or invalid.
+    """
     if not (source := body.get("_source")):
         return {}
 
@@ -28,6 +44,14 @@ def get_source_from_body(body: dict[str, Any]) -> dict[str, Any]:
 
 
 def get_aggregations_from_body(body: dict[str, Any]) -> dict[str, Any]:
+    """Extract and validate the aggregations field from an Elasticsearch response body.
+
+    Args:
+        body: The response body dictionary from Elasticsearch.
+
+    Returns:
+        The aggregations field as a dictionary, or an empty dict if missing or invalid.
+    """
     if not (aggregations := body.get("aggregations")):
         return {}
 
@@ -38,6 +62,17 @@ def get_aggregations_from_body(body: dict[str, Any]) -> dict[str, Any]:
 
 
 def get_hits_from_response(response: ObjectApiResponse[Any]) -> list[dict[str, Any]]:
+    """Extract and validate the hits array from an Elasticsearch response.
+
+    This function navigates the nested structure of Elasticsearch responses to extract
+    the hits.hits array which contains the actual search results.
+
+    Args:
+        response: The Elasticsearch API response object.
+
+    Returns:
+        A list of hit dictionaries from the response, or an empty list if the hits are missing or invalid.
+    """
     if not (body := response.body):  # pyright: ignore[reportAny]
         return []
 
@@ -66,6 +101,21 @@ def get_hits_from_response(response: ObjectApiResponse[Any]) -> list[dict[str, A
 
 
 def get_fields_from_hit(hit: dict[str, Any]) -> dict[str, list[Any]]:
+    """Extract and validate the fields from an Elasticsearch hit.
+
+    Elasticsearch can return stored fields via the "fields" key in each hit.
+    This function validates that the fields object exists and conforms to the
+    expected structure (a dict mapping field names to lists of values).
+
+    Args:
+        hit: A single hit dictionary from an Elasticsearch response.
+
+    Returns:
+        The fields dictionary from the hit, or an empty dict if missing.
+
+    Raises:
+        TypeError: If the fields structure is invalid (not a dict or contains non-list values).
+    """
     if not (fields := hit.get("fields")):
         return {}
 
@@ -81,6 +131,18 @@ def get_fields_from_hit(hit: dict[str, Any]) -> dict[str, list[Any]]:
 
 
 def get_field_from_hit(hit: dict[str, Any], field: str) -> list[Any]:
+    """Extract a specific field value from an Elasticsearch hit.
+
+    Args:
+        hit: A single hit dictionary from an Elasticsearch response.
+        field: The name of the field to extract.
+
+    Returns:
+        The field value as a list, or an empty list if the fields object is missing.
+
+    Raises:
+        TypeError: If the specified field is not present in the hit.
+    """
     if not (fields := get_fields_from_hit(hit=hit)):
         return []
 
@@ -92,6 +154,19 @@ def get_field_from_hit(hit: dict[str, Any], field: str) -> list[Any]:
 
 
 def get_values_from_field_in_hit(hit: dict[str, Any], field: str, value_type: type[T]) -> list[T]:
+    """Extract and type-check field values from an Elasticsearch hit.
+
+    Args:
+        hit: A single hit dictionary from an Elasticsearch response.
+        field: The name of the field to extract.
+        value_type: The expected type of values in the field list.
+
+    Returns:
+        A list of values of the specified type.
+
+    Raises:
+        TypeError: If the field is missing or contains values of the wrong type.
+    """
     if not (value := get_field_from_hit(hit=hit, field=field)):
         msg = f"Field {field} is not in hit {hit}"
         raise TypeError(msg)
@@ -104,6 +179,19 @@ def get_values_from_field_in_hit(hit: dict[str, Any], field: str, value_type: ty
 
 
 def get_first_value_from_field_in_hit(hit: dict[str, Any], field: str, value_type: type[T]) -> T:
+    """Extract and validate a single-value field from an Elasticsearch hit.
+
+    Args:
+        hit: A single hit dictionary from an Elasticsearch response.
+        field: The name of the field to extract.
+        value_type: The expected type of the value.
+
+    Returns:
+        The single value from the field.
+
+    Raises:
+        TypeError: If the field doesn't contain exactly one value, or if the value type is incorrect.
+    """
     values: list[T] = get_values_from_field_in_hit(hit=hit, field=field, value_type=value_type)
     if len(values) != 1:
         msg: str = f"Field {field} in hit {hit} is not a single value"
@@ -112,6 +200,19 @@ def get_first_value_from_field_in_hit(hit: dict[str, Any], field: str, value_typ
 
 
 def managed_entry_to_document(collection: str, key: str, managed_entry: ManagedEntry) -> dict[str, Any]:
+    """Convert a ManagedEntry to an Elasticsearch document format.
+
+    This function serializes a ManagedEntry to a document suitable for storage in Elasticsearch,
+    including collection, key, and value fields, plus optional timestamp metadata.
+
+    Args:
+        collection: The collection name to include in the document.
+        key: The key to include in the document.
+        managed_entry: The ManagedEntry to serialize.
+
+    Returns:
+        An Elasticsearch document dictionary ready for indexing.
+    """
     document: dict[str, Any] = {"collection": collection, "key": key, "value": managed_entry.to_json(include_metadata=False)}
 
     if managed_entry.created_at:
@@ -123,4 +224,14 @@ def managed_entry_to_document(collection: str, key: str, managed_entry: ManagedE
 
 
 def new_bulk_action(action: str, index: str, document_id: str) -> dict[str, Any]:
+    """Create a bulk action descriptor for Elasticsearch bulk operations.
+
+    Args:
+        action: The bulk action type (e.g., "index", "delete", "update").
+        index: The Elasticsearch index name.
+        document_id: The document ID.
+
+    Returns:
+        A bulk action dictionary formatted for Elasticsearch's bulk API.
+    """
     return {action: {"_index": index, "_id": document_id}}