pyrmute 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

pyrmute/__init__.py

@@ -1,20 +1,47 @@
-"""pyrmute - versioned Pydantic models and schemas with migrations.
-
-A package for managing versioned Pydantic models with automatic migrations
-and schema management.
-"""
+"""pyrmute - versioned Pydantic models and schemas with migrations."""
 
+from ._migration_manager import MigrationManager
+from ._registry import Registry
+from ._schema_manager import SchemaManager
 from ._version import __version__
+from .exceptions import (
+    InvalidVersionError,
+    MigrationError,
+    ModelNotFoundError,
+    VersionedModelError,
+)
+from .migration_testing import (
+    MigrationTestCase,
+    MigrationTestResult,
+    MigrationTestResults,
+)
+from .model_diff import ModelDiff
 from .model_manager import ModelManager
 from .model_version import ModelVersion
-from .types import JsonSchema, MigrationData, MigrationFunc, ModelMetadata
+from .types import (
+    JsonSchema,
+    MigrationFunc,
+    ModelData,
+    NestedModelInfo,
+)
 
 __all__ = [
+    "InvalidVersionError",
     "JsonSchema",
-    "MigrationData",
+    "MigrationError",
     "MigrationFunc",
+    "MigrationManager",
+    "MigrationTestCase",
+    "MigrationTestResult",
+    "MigrationTestResults",
+    "ModelData",
+    "ModelDiff",
     "ModelManager",
-    "ModelMetadata",
+    "ModelNotFoundError",
     "ModelVersion",
+    "NestedModelInfo",
+    "Registry",
+    "SchemaManager",
+    "VersionedModelError",
     "__version__",
 ]

pyrmute/_migration_manager.py

@@ -1,24 +1,27 @@
 """Migrations manager."""
 
+import contextlib
 from collections.abc import Callable
 from typing import Any, Self, get_args, get_origin
 
 from pydantic import BaseModel
 from pydantic.fields import FieldInfo
+from pydantic_core import PydanticUndefined
 
 from ._registry import Registry
+from .exceptions import MigrationError, ModelNotFoundError
 from .model_version import ModelVersion
-from .types import MigrationData, MigrationFunc, ModelName
+from .types import MigrationFunc, ModelData, ModelName
 
 
 class MigrationManager:
     """Manager for data migrations between model versions.
 
-    Handles registration and execution of migration functions, including
-    support for nested Pydantic models.
+    Handles registration and execution of migration functions, including support for
+    nested Pydantic models.
 
     Attributes:
-        registry: Reference to the VersionedModelRegistry.
+        registry: Reference to the Registry.
     """
 
     def __init__(self: Self, registry: Registry) -> None:
@@ -70,11 +73,11 @@ class MigrationManager:
 
     def migrate(
         self: Self,
-        data: MigrationData,
+        data: ModelData,
         name: ModelName,
         from_version: str | ModelVersion,
         to_version: str | ModelVersion,
-    ) -> MigrationData:
+    ) -> ModelData:
         """Migrate data from one version to another.
 
         Args:
@@ -87,7 +90,8 @@ class MigrationManager:
             Migrated data dictionary.
 
         Raises:
-            ValueError: If migration path cannot be found.
+            ModelNotFoundError: If model or versions don't exist.
+            MigrationError: If migration path cannot be found.
         """
         from_ver = (
             ModelVersion.parse(from_version)
@@ -103,7 +107,7 @@ class MigrationManager:
         if from_ver == to_ver:
             return data
 
-        path = self._find_migration_path(name, from_ver, to_ver)
+        path = self.find_migration_path(name, from_ver, to_ver)
 
         current_data = data
         for i in range(len(path) - 1):
@@ -111,15 +115,41 @@ class MigrationManager:
 
             if migration_key in self.registry._migrations[name]:
                 migration_func = self.registry._migrations[name][migration_key]
-                current_data = migration_func(current_data)
+                try:
+                    current_data = migration_func(current_data)
+                except Exception as e:
+                    raise MigrationError(
+                        name,
+                        str(path[i]),
+                        str(path[i + 1]),
+                        f"Migration function raised: {type(e).__name__}: {e}",
+                    ) from e
+            elif path[i + 1] in self.registry._backward_compatible_enabled[name]:
+                try:
+                    current_data = self._auto_migrate(
+                        current_data, name, path[i], path[i + 1]
+                    )
+                except Exception as e:
+                    raise MigrationError(
+                        name,
+                        str(path[i]),
+                        str(path[i + 1]),
+                        f"Auto-migration failed: {type(e).__name__}: {e}",
+                    ) from e
             else:
-                current_data = self._auto_migrate(
-                    current_data, name, path[i], path[i + 1]
+                raise MigrationError(
+                    name,
+                    str(path[i]),
+                    str(path[i + 1]),
+                    (
+                        "No migration path found. Define a migration function or mark "
+                        "the target version as backward_compatible."
+                    ),
                 )
 
         return current_data
 
-    def _find_migration_path(
+    def find_migration_path(
         self: Self,
         name: ModelName,
         from_ver: ModelVersion,
@@ -134,11 +164,16 @@ class MigrationManager:
 
         Returns:
             List of versions forming the migration path.
+
+        Raises:
+            ModelNotFoundError: If the model or versions don't exist.
         """
         versions = sorted(self.registry.get_versions(name))
 
-        if from_ver not in versions or to_ver not in versions:
-            raise ValueError(f"Invalid version range for {name}")
+        if from_ver not in versions:
+            raise ModelNotFoundError(name, str(from_ver))
+        if to_ver not in versions:
+            raise ModelNotFoundError(name, str(to_ver))
 
         from_idx = versions.index(from_ver)
         to_idx = versions.index(to_ver)
@@ -147,17 +182,57 @@ class MigrationManager:
             return versions[from_idx : to_idx + 1]
         return versions[to_idx : from_idx + 1][::-1]
 
+    def validate_migration_path(
+        self: Self,
+        name: ModelName,
+        from_ver: ModelVersion,
+        to_ver: ModelVersion,
+    ) -> None:
+        """Validate that a migration path exists and all steps are valid.
+
+        Args:
+            name: Name of the model.
+            from_ver: Source version.
+            to_ver: Target version.
+
+        Raises:
+            ModelNotFoundError: If the model or versions don't exist.
+            MigrationError: If any step in the migration path is invalid.
+        """
+        path = self.find_migration_path(name, from_ver, to_ver)
+
+        for i in range(len(path) - 1):
+            current_ver = path[i]
+            next_ver = path[i + 1]
+            migration_key = (current_ver, next_ver)
+
+            has_explicit = migration_key in self.registry._migrations.get(name, {})
+            has_auto = next_ver in self.registry._backward_compatible_enabled.get(
+                name, set()
+            )
+
+            if not has_explicit and not has_auto:
+                raise MigrationError(
+                    name,
+                    str(current_ver),
+                    str(next_ver),
+                    (
+                        "No migration path found. Define a migration function or mark "
+                        "the target version as backward_compatible."
+                    ),
+                )
+
     def _auto_migrate(
         self: Self,
-        data: MigrationData,
+        data: ModelData,
         name: ModelName,
         from_ver: ModelVersion,
         to_ver: ModelVersion,
-    ) -> MigrationData:
+    ) -> ModelData:
         """Automatically migrate data when no explicit migration exists.
 
-        This method handles nested Pydantic models recursively, migrating
-        them to their corresponding versions.
+        This method handles nested Pydantic models recursively, migrating them to their
+        corresponding versions.
 
         Args:
             data: Data dictionary to migrate.
@@ -174,21 +249,28 @@ class MigrationManager:
         from_fields = from_model.model_fields
         to_fields = to_model.model_fields
 
-        result: MigrationData = {}
+        result: ModelData = {}
 
         for field_name, to_field_info in to_fields.items():
-            if field_name not in data:
-                continue
-
-            value = data[field_name]
+            # Field exists in data, migrate it
+            if field_name in data:
+                value = data[field_name]
+                from_field_info = from_fields.get(field_name)
+                result[field_name] = self._migrate_field_value(
+                    value, from_field_info, to_field_info
+                )
 
-            # Get corresponding from_field if it exists
-            from_field_info = from_fields.get(field_name)
+            # Field missing from data, use default if available
+            elif to_field_info.default is not PydanticUndefined:
+                result[field_name] = to_field_info.default
+            elif to_field_info.default_factory is not None:
+                with contextlib.suppress(Exception):
+                    result[field_name] = to_field_info.default_factory()  # type: ignore
 
-            # Migrate the field value (handles nested models)
-            result[field_name] = self._migrate_field_value(
-                value, from_field_info, to_field_info
-            )
+        # Migrate all extra data not in the field, too
+        for field_name, value in data.items():
+            if field_name not in to_fields:
+                result[field_name] = value
 
         return result
 
@@ -211,20 +293,17 @@ class MigrationManager:
         if value is None:
             return None
 
-        # Check if this is a nested Pydantic model
         if isinstance(value, dict):
             nested_info = self._extract_nested_model_info(value, from_field, to_field)
             if nested_info:
                 nested_name, nested_from_ver, nested_to_ver = nested_info
                 return self.migrate(value, nested_name, nested_from_ver, nested_to_ver)
 
-            # Try to recursively migrate dict values
             return {
                 k: self._migrate_field_value(v, from_field, to_field)
                 for k, v in value.items()
             }
 
-        # Handle lists
         if isinstance(value, list):
             return [
                 self._migrate_field_value(item, from_field, to_field) for item in value
@@ -234,7 +313,7 @@ class MigrationManager:
 
     def _extract_nested_model_info(
         self: Self,
-        value: MigrationData,
+        value: ModelData,
         from_field: FieldInfo | None,
         to_field: FieldInfo,
     ) -> tuple[ModelName, ModelVersion, ModelVersion] | None:
@@ -249,12 +328,10 @@ class MigrationManager:
             Tuple of (model_name, from_version, to_version) if this is a
             versioned nested model, None otherwise.
         """
-        # Get the target model type
         to_model_type = self._get_model_type_from_field(to_field)
         if not to_model_type or not issubclass(to_model_type, BaseModel):
             return None
 
-        # Check if target model is registered
         to_info = self.registry.get_model_info(to_model_type)
         if not to_info:
             return None
@@ -291,11 +368,9 @@ class MigrationManager:
         if annotation is None:
             return None
 
-        # Handle direct model types
         if isinstance(annotation, type) and issubclass(annotation, BaseModel):
             return annotation
 
-        # Handle Optional, List, etc.
         origin = get_origin(annotation)
         if origin is not None:
             args = get_args(annotation)

pyrmute/_registry.py

@@ -6,6 +6,7 @@ from typing import Self
 
 from pydantic import BaseModel
 
+from .exceptions import ModelNotFoundError
 from .model_version import ModelVersion
 from .types import (
     DecoratedBaseModel,
@@ -21,8 +22,8 @@ from .types import (
 class Registry:
     """Registry for versioned Pydantic models.
 
-    Manages the registration and retrieval of versioned models and their
-    associated metadata.
+    Manages the registration and retrieval of versioned models and their associated
+    metadata.
 
     Attributes:
         _models: Dictionary mapping model names to version-model mappings.
@@ -39,6 +40,9 @@ class Registry:
         self._schema_generators: dict[ModelName, SchemaGenerators] = defaultdict(dict)
         self._model_metadata: dict[type[BaseModel], ModelMetadata] = {}
         self._ref_enabled: dict[ModelName, set[ModelVersion]] = defaultdict(set)
+        self._backward_compatible_enabled: dict[ModelName, set[ModelVersion]] = (
+            defaultdict(set)
+        )
 
     def register(
         self: Self,
@@ -46,6 +50,7 @@ class Registry:
         version: str | ModelVersion,
         schema_generator: JsonSchemaGenerator | None = None,
         enable_ref: bool = False,
+        backward_compatible: bool = False,
     ) -> Callable[[type[DecoratedBaseModel]], type[DecoratedBaseModel]]:
         """Register a versioned model.
 
@@ -53,8 +58,11 @@ class Registry:
             name: Name of the model.
             version: Semantic version string or ModelVersion instance.
             schema_generator: Optional custom schema generator function.
-            enable_ref: If True, this model can be referenced via $ref in
-                separate schema files. If False, it will always be inlined.
+            enable_ref: If True, this model can be referenced via $ref in separate
+                schema files. If False, it will always be inlined.
+            backward_compatible: If True, this model does not need a migration function
+                to migrate to the next version. If a migration function is defined it
+                will use it.
 
         Returns:
             Decorator function for model class.
@@ -74,6 +82,8 @@ class Registry:
                 self._schema_generators[name][ver] = schema_generator
             if enable_ref:
                 self._ref_enabled[name].add(ver)
+            if backward_compatible:
+                self._backward_compatible_enabled[name].add(ver)
             return cls
 
         return decorator
@@ -91,12 +101,15 @@ class Registry:
             Model class for the specified version.
 
         Raises:
-            ValueError: If model or version not found.
+            ModelNotFoundError: If model or version not found.
         """
         ver = ModelVersion.parse(version) if isinstance(version, str) else version
 
-        if name not in self._models or ver not in self._models[name]:
-            raise ValueError(f"Model {name} v{ver} not found")
+        if name not in self._models:
+            raise ModelNotFoundError(name)
+
+        if ver not in self._models[name]:
+            raise ModelNotFoundError(name, str(ver))
 
         return self._models[name][ver]
 
@@ -110,10 +123,10 @@ class Registry:
             Latest version of the model class.
 
         Raises:
-            ValueError: If model not found.
+            ModelNotFoundError: If model not found.
         """
         if name not in self._models:
-            raise ValueError(f"Model {name} not found")
+            raise ModelNotFoundError(name)
 
         latest_version = max(self._models[name].keys())
         return self._models[name][latest_version]
@@ -128,10 +141,10 @@ class Registry:
             Sorted list of available versions.
 
         Raises:
-            ValueError: If model not found.
+            ModelNotFoundError: If model not found.
         """
         if name not in self._models:
-            raise ValueError(f"Model {name} not found")
+            raise ModelNotFoundError(name)
 
         return sorted(self._models[name].keys())
 

pyrmute/_schema_manager.py

@@ -8,6 +8,7 @@ from pydantic import BaseModel
 from pydantic.fields import FieldInfo
 
 from ._registry import Registry
+from .exceptions import ModelNotFoundError
 from .model_version import ModelVersion
 from .types import (
     JsonSchema,
@@ -15,14 +16,15 @@ from .types import (
     JsonValue,
     ModelMetadata,
     ModelName,
+    NestedModelInfo,
 )
 
 
 class SchemaManager:
     """Manager for JSON schema generation and export.
 
-    Handles schema generation from Pydantic models with support for
-    custom schema generators and sub-schema references.
+    Handles schema generation from Pydantic models with support for custom schema
+    generators and sub-schema references.
 
     Attributes:
         registry: Reference to the Registry.
@@ -73,14 +75,14 @@ class SchemaManager:
     ) -> JsonSchema:
         """Get JSON schema with separate definition files for nested models.
 
-        This creates a schema where nested Pydantic models are referenced
-        as external JSON schema files rather than inline definitions.
+        This creates a schema where nested Pydantic models are referenced as external
+        JSON schema files rather than inline definitions.
 
         Args:
             name: Name of the model.
             version: Semantic version.
-            ref_template: Template for generating $ref URLs. Supports {model}
-                and {version} placeholders.
+            ref_template: Template for generating $ref URLs. Supports {model} and
+                {version} placeholders.
             **schema_kwargs: Additional arguments for schema generation.
 
         Returns:
@@ -93,8 +95,6 @@ class SchemaManager:
             ... )
         """
         ver = ModelVersion.parse(version) if isinstance(version, str) else version
-
-        # Get the base schema with definitions
         schema = self.get_schema(name, ver, **schema_kwargs)
 
         # Extract and replace definitions with external references
@@ -136,16 +136,15 @@ class SchemaManager:
                 if "$ref" in value:
                     # Extract the definition name from the ref
                     ref = value["$ref"]
-                    if ref.startswith(("#/$defs/", "#/definitions/")):
+                    if isinstance(ref, str) and ref.startswith(
+                        ("#/$defs/", "#/definitions/")
+                    ):
                         def_name = ref.split("/")[-1]
 
-                        # Try to find the model info for this definition
                         model_info = self._find_model_for_definition(def_name)
-
                         if model_info:
                             model_name, model_version = model_info
 
-                            # Check if this model is enabled for $ref
                             if self.registry.is_ref_enabled(model_name, model_version):
                                 # Replace with external reference
                                 return {
@@ -156,7 +155,6 @@ class SchemaManager:
                             # Keep as internal reference (will be inlined)
                             return value
 
-                # Recursively process nested dictionaries
                 return {k: process_value(v) for k, v in value.items()}
             if isinstance(value, list):
                 return [process_value(item) for item in value]
@@ -178,7 +176,6 @@ class SchemaManager:
         Returns:
             Dictionary of definitions that weren't converted to external refs.
         """
-        # Find all internal refs still in the schema
         internal_refs: set[str] = set()
 
         def find_internal_refs(value: dict[str, Any] | list[Any]) -> None:
@@ -195,8 +192,6 @@ class SchemaManager:
                     find_internal_refs(item)
 
         find_internal_refs(schema)
-
-        # Return only definitions that are still referenced internally
         return {k: v for k, v in original_defs.items() if k in internal_refs}
 
     def _find_model_for_definition(self: Self, def_name: str) -> ModelMetadata | None:
@@ -208,7 +203,6 @@ class SchemaManager:
         Returns:
             Tuple of (model_name, version) if found, None otherwise.
         """
-        # Search through all registered models to find matching class name
         for name, versions in self.registry._models.items():
             for version, model_class in versions.items():
                 if model_class.__name__ == def_name:
@@ -225,10 +219,10 @@ class SchemaManager:
             Dictionary mapping versions to their schemas.
 
         Raises:
-            ValueError: If model not found.
+            ModelNotFoundError: If model not found.
         """
         if name not in self.registry._models:
-            raise ValueError(f"Model {name} not found")
+            raise ModelNotFoundError(name)
 
         return {
             version: self.get_schema(name, version)
@@ -247,8 +241,8 @@ class SchemaManager:
         Args:
             output_dir: Directory path for output files.
             indent: JSON indentation level.
-            separate_definitions: If True, create separate schema files for
-                nested models that have enable_ref=True.
+            separate_definitions: If True, create separate schema files for nested
+                models that have enable_ref=True.
             ref_template: Template for $ref URLs when separate_definitions=True.
                 Defaults to relative file references if not provided.
 
@@ -270,19 +264,15 @@ class SchemaManager:
         output_path.mkdir(parents=True, exist_ok=True)
 
         if not separate_definitions:
-            # Original behavior: inline definitions
             for name in self.registry._models:
                 for version, schema in self.get_all_schemas(name).items():
                     file_path = output_path / f"{name}_v{version}.json"
                     with open(file_path, "w", encoding="utf-8") as f:
                         json.dump(schema, f, indent=indent)
         else:
-            # New behavior: separate definition files for enable_ref=True models
-            # Default to relative file references
             if ref_template is None:
                 ref_template = "{model}_v{version}.json"
 
-            # First pass: write all schemas
             for name in self.registry._models:
                 for version in self.registry._models[name]:
                     schema = self.get_schema_with_separate_defs(
@@ -296,7 +286,7 @@ class SchemaManager:
         self: Self,
         name: ModelName,
         version: str | ModelVersion,
-    ) -> list[ModelMetadata]:
+    ) -> list[NestedModelInfo]:
         """Get all nested models referenced by a model.
 
         Args:
@@ -304,19 +294,28 @@ class SchemaManager:
             version: Semantic version.
 
         Returns:
-            List of (model_name, model_version) tuples for nested models.
+            List of NestedModelInfo.
         """
         ver = ModelVersion.parse(version) if isinstance(version, str) else version
         model = self.registry.get_model(name, ver)
 
-        nested: list[ModelMetadata] = []
+        nested: list[NestedModelInfo] = []
 
         for field_info in model.model_fields.values():
             model_type = self._get_model_type_from_field(field_info)
-            if model_type:
-                model_info = self.registry.get_model_info(model_type)
-                if model_info and model_info not in nested:
-                    nested.append(model_info)
+            if not model_type:
+                continue
+
+            model_info = self.registry.get_model_info(model_type)
+
+            if not model_info:
+                continue
+
+            name_, version_ = model_info
+            nested_model_info = NestedModelInfo(name=name_, version=version_)
+
+            if nested_model_info not in nested:
+                nested.append(nested_model_info)
 
         return nested
 
@@ -335,11 +334,9 @@ class SchemaManager:
         if annotation is None:
             return None
 
-        # Handle direct model types
         if isinstance(annotation, type) and issubclass(annotation, BaseModel):
             return annotation
 
-        # Handle Optional, List, etc.
         origin = get_origin(annotation)
         if origin is not None:
             args = get_args(annotation)

pyrmute/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.0'
-__version_tuple__ = version_tuple = (0, 1, 0)
+__version__ = version = '0.3.0'
+__version_tuple__ = version_tuple = (0, 3, 0)
 
 __commit_id__ = commit_id = None