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

pyrmute/__init__.py

@@ -12,6 +12,7 @@ from .exceptions import (
 )
 from .migration_testing import (
     MigrationTestCase,
+    MigrationTestCases,
     MigrationTestResult,
     MigrationTestResults,
 )
@@ -32,6 +33,7 @@ __all__ = [
     "MigrationFunc",
     "MigrationManager",
     "MigrationTestCase",
+    "MigrationTestCases",
     "MigrationTestResult",
     "MigrationTestResults",
     "ModelData",

pyrmute/_migration_manager.py

@@ -1,8 +1,9 @@
 """Migrations manager."""
 
 import contextlib
+import types
 from collections.abc import Callable
-from typing import Any, Self, get_args, get_origin
+from typing import Annotated, Any, Literal, Self, Union, get_args, get_origin
 
 from pydantic import BaseModel
 from pydantic.fields import FieldInfo
@@ -54,16 +55,8 @@ class MigrationManager:
             ... def migrate_v1_to_v2(data: dict[str, Any]) -> dict[str, Any]:
             ...     return {**data, "email": "unknown@example.com"}
         """
-        from_ver = (
-            ModelVersion.parse(from_version)
-            if isinstance(from_version, str)
-            else from_version
-        )
-        to_ver = (
-            ModelVersion.parse(to_version)
-            if isinstance(to_version, str)
-            else to_version
-        )
+        from_ver = self._parse_version(from_version)
+        to_ver = self._parse_version(to_version)
 
         def decorator(func: MigrationFunc) -> MigrationFunc:
             self.registry._migrations[name][(from_ver, to_ver)] = func
@@ -93,61 +86,14 @@ class MigrationManager:
             ModelNotFoundError: If model or versions don't exist.
             MigrationError: If migration path cannot be found.
         """
-        from_ver = (
-            ModelVersion.parse(from_version)
-            if isinstance(from_version, str)
-            else from_version
-        )
-        to_ver = (
-            ModelVersion.parse(to_version)
-            if isinstance(to_version, str)
-            else to_version
-        )
+        from_ver = self._parse_version(from_version)
+        to_ver = self._parse_version(to_version)
 
         if from_ver == to_ver:
             return data
 
         path = self.find_migration_path(name, from_ver, to_ver)
-
-        current_data = data
-        for i in range(len(path) - 1):
-            migration_key = (path[i], path[i + 1])
-
-            if migration_key in self.registry._migrations[name]:
-                migration_func = self.registry._migrations[name][migration_key]
-                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:
-                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
+        return self._execute_migration_path(data, name, path)
 
     def find_migration_path(
         self: Self,
@@ -168,6 +114,9 @@ class MigrationManager:
         Raises:
             ModelNotFoundError: If the model or versions don't exist.
         """
+        if (from_ver, to_ver) in self.registry._migrations.get(name, {}):
+            return [from_ver, to_ver]
+
         versions = sorted(self.registry.get_versions(name))
 
         if from_ver not in versions:
@@ -204,14 +153,8 @@ class MigrationManager:
         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:
+            if not self._has_migration_step(name, current_ver, next_ver):
                 raise MigrationError(
                     name,
                     str(current_ver),
@@ -222,6 +165,85 @@ class MigrationManager:
                     ),
                 )
 
+    def _parse_version(self: Self, version: str | ModelVersion) -> ModelVersion:
+        """Parse version string or return ModelVersion as-is."""
+        return ModelVersion.parse(version) if isinstance(version, str) else version
+
+    def _has_migration_step(
+        self: Self, name: ModelName, from_ver: ModelVersion, to_ver: ModelVersion
+    ) -> bool:
+        """Check if a migration step exists (explicit or auto)."""
+        migration_key = (from_ver, to_ver)
+        has_explicit = migration_key in self.registry._migrations.get(name, {})
+        has_auto = to_ver in self.registry._backward_compatible_enabled.get(name, set())
+        return has_explicit or has_auto
+
+    def _execute_migration_path(
+        self: Self, data: ModelData, name: ModelName, path: list[ModelVersion]
+    ) -> ModelData:
+        """Execute migration through a path of versions."""
+        current_data = data
+
+        for i in range(len(path) - 1):
+            try:
+                current_data = self._execute_single_step(
+                    current_data, name, path[i], path[i + 1]
+                )
+            except Exception as e:
+                if isinstance(e, MigrationError):
+                    raise
+                raise MigrationError(
+                    name,
+                    str(path[i]),
+                    str(path[i + 1]),
+                    f"Migration failed: {type(e).__name__}: {e}",
+                ) from e
+
+        return current_data
+
+    def _execute_single_step(
+        self: Self,
+        data: ModelData,
+        name: ModelName,
+        from_ver: ModelVersion,
+        to_ver: ModelVersion,
+    ) -> ModelData:
+        """Execute a single migration step."""
+        migration_key = (from_ver, to_ver)
+
+        if migration_key in self.registry._migrations[name]:
+            migration_func = self.registry._migrations[name][migration_key]
+            try:
+                return migration_func(data)
+            except Exception as e:
+                raise MigrationError(
+                    name,
+                    str(from_ver),
+                    str(to_ver),
+                    f"Migration function raised: {type(e).__name__}: {e}",
+                ) from e
+
+        if to_ver in self.registry._backward_compatible_enabled[name]:
+            try:
+                return self._auto_migrate(data, name, from_ver, to_ver)
+            except Exception as e:
+                raise MigrationError(
+                    name,
+                    str(from_ver),
+                    str(to_ver),
+                    f"Auto-migration failed: {type(e).__name__}: {e}",
+                ) from e
+
+        raise MigrationError(
+            name,
+            str(from_ver),
+            str(to_ver),
+            (
+                "No migration path found. Define a migration function or mark "
+                "the target version as backward_compatible."
+            ),
+        )
+
     def _auto_migrate(
         self: Self,
         data: ModelData,
@@ -232,7 +254,7 @@ class MigrationManager:
         """Automatically migrate data when no explicit migration exists.
 
         This method handles nested Pydantic models recursively, migrating them to their
-        corresponding versions.
+        corresponding versions. Handles field aliases by building a lookup map.
 
         Args:
             data: Data dictionary to migrate.
@@ -249,36 +271,142 @@ class MigrationManager:
         from_fields = from_model.model_fields
         to_fields = to_model.model_fields
 
+        key_to_field_name = self._build_alias_map(from_fields)
+
         result: ModelData = {}
+        processed_keys: set[str] = set()
 
         for field_name, to_field_info in to_fields.items():
-            # 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
-                )
+            field_result = self._migrate_single_field(
+                data, field_name, to_field_info, from_fields, key_to_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
+            if field_result is not None:
+                output_key, migrated_value, source_keys = field_result
+                result[output_key] = migrated_value
+                processed_keys.update(source_keys)
 
-        # 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
+        # Preserve unprocessed extra fields
+        for data_key, value in data.items():
+            if data_key not in processed_keys:
+                result[data_key] = value
 
         return result
 
-    def _migrate_field_value(
+    def _build_alias_map(self: Self, fields: dict[str, FieldInfo]) -> dict[str, str]:
+        """Build a mapping from all possible input keys to canonical field names.
+
+        Args:
+            fields: Model fields to extract aliases from.
+
+        Returns:
+            Dictionary mapping data keys (field names and aliases) to field names.
+        """
+        key_to_field_name: dict[str, str] = {}
+
+        for field_name, field_info in fields.items():
+            key_to_field_name[field_name] = field_name
+
+            if field_info.alias:
+                key_to_field_name[field_info.alias] = field_name
+            if field_info.serialization_alias:
+                key_to_field_name[field_info.serialization_alias] = field_name
+            if isinstance(field_info.validation_alias, str):
+                key_to_field_name[field_info.validation_alias] = field_name
+
+        return key_to_field_name
+
+    def _migrate_single_field(
         self: Self,
-        value: Any,
-        from_field: FieldInfo | None,
-        to_field: FieldInfo,
+        data: ModelData,
+        field_name: str,
+        to_field_info: FieldInfo,
+        from_fields: dict[str, FieldInfo],
+        key_to_field_name: dict[str, str],
+    ) -> tuple[str, Any, set[str]] | None:
+        """Migrate a single field, handling aliases and defaults.
+
+        Args:
+            data: Source data dictionary.
+            field_name: Target field name.
+            to_field_info: Target field info.
+            from_fields: Source model fields.
+            key_to_field_name: Mapping from data keys to field names.
+
+        Returns:
+            Tuple of (output_key, migrated_value, source_keys_to_mark_processed) if
+            field should be included, None if field should be skipped.
+        """
+        value, data_key_used = self._find_field_value(
+            data, field_name, key_to_field_name
+        )
+
+        if data_key_used is not None:
+            from_field_info = from_fields.get(field_name)
+            migrated_value = self._migrate_field_value(
+                value, from_field_info, to_field_info
+            )
+            keys_to_process = {
+                k for k, v in key_to_field_name.items() if v == field_name
+            }
+            return (data_key_used, migrated_value, keys_to_process)
+
+        _NO_DEFAULT = object()
+        default_value = self._get_field_default(to_field_info, _NO_DEFAULT)
+
+        if default_value is not _NO_DEFAULT:
+            output_key = (
+                to_field_info.serialization_alias or to_field_info.alias or field_name
+            )
+            return (output_key, default_value, set())
+
+        return None
+
+    def _find_field_value(
+        self: Self, data: ModelData, field_name: str, key_to_field_name: dict[str, str]
+    ) -> tuple[Any, str | None]:
+        """Find a field's value in data, checking both field name and aliases.
+
+        Args:
+            data: Source data dictionary.
+            field_name: Target field name to find.
+            key_to_field_name: Mapping from data keys to field names.
+
+        Returns:
+            Tuple of (value, data_key) if found, (None, None) otherwise.
+        """
+        if field_name in data:
+            return (data[field_name], field_name)
+
+        for data_key in data:
+            if key_to_field_name.get(data_key) == field_name:
+                return (data[data_key], data_key)
+
+        return (None, None)
+
+    def _get_field_default(
+        self: Self, field_info: FieldInfo, sentinel: Any = None
+    ) -> Any:
+        """Get the default value for a field.
+
+        Args:
+            field_info: Field info to extract default from.
+            sentinel: Sentinel value to return if no default is available.
+
+        Returns:
+            Default value if available, sentinel otherwise.
+        """
+        if field_info.default is not PydanticUndefined:
+            return field_info.default
+
+        if field_info.default_factory is not None:
+            with contextlib.suppress(Exception):
+                return field_info.default_factory()  # type: ignore
+
+        return sentinel
+
+    def _migrate_field_value(
+        self: Self, value: Any, from_field: FieldInfo | None, to_field: FieldInfo
     ) -> Any:
         """Migrate a single field value, handling nested models.
 
@@ -294,31 +422,104 @@ class MigrationManager:
             return None
 
         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)
-
-            return {
-                k: self._migrate_field_value(v, from_field, to_field)
-                for k, v in value.items()
-            }
+            return self._migrate_dict_value(value, from_field, to_field)
 
         if isinstance(value, list):
-            return [
-                self._migrate_field_value(item, from_field, to_field) for item in value
-            ]
+            return self._migrate_list_value(value, from_field, to_field)
 
         return value
 
+    def _migrate_dict_value(
+        self, value: dict[str, Any], from_field: FieldInfo | None, to_field: FieldInfo
+    ) -> dict[str, Any]:
+        """Migrate a dictionary value (might be a nested model)."""
+        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)
+
+        return {
+            k: self._migrate_field_value(v, from_field, to_field)
+            for k, v in value.items()
+        }
+
+    def _migrate_list_value(
+        self, value: list[Any], from_field: FieldInfo | None, to_field: FieldInfo
+    ) -> list[Any]:
+        """Migrate a list value (might contain nested models)."""
+        from_item_field, to_item_field = self._get_list_item_fields(
+            from_field, to_field
+        )
+
+        return [
+            self._migrate_list_item(item, from_item_field, to_item_field)
+            for item in value
+        ]
+
+    def _migrate_list_item(
+        self, item: Any, from_field: FieldInfo | None, to_field: FieldInfo | None
+    ) -> Any:
+        """Migrate a single item from a list."""
+        if not isinstance(item, dict) or to_field is None:
+            return item
+
+        nested_info = self._extract_nested_model_info(item, from_field, to_field)
+        if nested_info:
+            nested_name, nested_from_ver, nested_to_ver = nested_info
+            return self.migrate(item, nested_name, nested_from_ver, nested_to_ver)
+
+        return item
+
+    def _get_list_item_fields(
+        self, from_field: FieldInfo | None, to_field: FieldInfo
+    ) -> tuple[FieldInfo | None, FieldInfo | None]:
+        """Extract field info for items in a list field."""
+        to_item_field = self._extract_list_item_field(to_field)
+        from_item_field = (
+            self._extract_list_item_field(from_field) if from_field else None
+        )
+        return from_item_field, to_item_field
+
+    def _extract_list_item_field(self, field: FieldInfo | None) -> FieldInfo | None:
+        """Extract field info for the items of a list field."""
+        if field is None or field.annotation is None:
+            return None
+
+        origin = get_origin(field.annotation)
+        if origin is not list:
+            return None
+
+        args = get_args(field.annotation)
+        if not args:
+            return None
+
+        item_annotation = args[0]
+
+        discriminator = None
+        if get_origin(item_annotation) is Annotated:
+            annotated_args = get_args(item_annotation)
+            item_annotation = annotated_args[0]
+            for metadata in annotated_args[1:]:
+                if hasattr(metadata, "discriminator"):
+                    discriminator = metadata.discriminator
+
+        synthetic_field = FieldInfo(
+            annotation=item_annotation, default=PydanticUndefined
+        )
+        if discriminator:
+            synthetic_field.discriminator = discriminator
+
+        return synthetic_field
+
     def _extract_nested_model_info(
-        self: Self,
-        value: ModelData,
-        from_field: FieldInfo | None,
-        to_field: FieldInfo,
+        self, value: ModelData, from_field: FieldInfo | None, to_field: FieldInfo
     ) -> tuple[ModelName, ModelVersion, ModelVersion] | None:
         """Extract nested model migration information.
 
+        Handles discriminated unions by using the discriminator field to determine which
+        model type to migrate.
+
+
         Args:
             value: The nested model data.
             from_field: Source field info.
@@ -328,8 +529,51 @@ class MigrationManager:
             Tuple of (model_name, from_version, to_version) if this is a
             versioned nested model, None otherwise.
         """
+        discriminated_info = self._try_extract_discriminated_model(
+            value, from_field, to_field
+        )
+        if discriminated_info:
+            return discriminated_info
+
+        return self._try_extract_simple_nested_model(from_field, to_field)
+
+    def _try_extract_discriminated_model(
+        self, value: ModelData, from_field: FieldInfo | None, to_field: FieldInfo
+    ) -> tuple[ModelName, ModelVersion, ModelVersion] | None:
+        """Try to extract model info from a discriminated union field."""
+        discriminator_key = self._get_discriminator_key(to_field)
+        if not discriminator_key:
+            return None
+
+        discriminator_value = self._find_discriminator_value(
+            value, discriminator_key, to_field
+        )
+        if discriminator_value is None:
+            return None
+
+        to_model_type = self._find_discriminated_type(
+            to_field, discriminator_key, discriminator_value
+        )
+        if not to_model_type:
+            return None
+
+        to_info = self.registry.get_model_info(to_model_type)
+        if not to_info:
+            return None
+
+        model_name, to_version = to_info
+        from_version = self._get_discriminated_source_version(
+            from_field, discriminator_key, discriminator_value, model_name, to_version
+        )
+
+        return (model_name, from_version, to_version)
+
+    def _try_extract_simple_nested_model(
+        self, from_field: FieldInfo | None, to_field: FieldInfo
+    ) -> tuple[ModelName, ModelVersion, ModelVersion] | None:
+        """Try to extract model info from a simple (non-discriminated) nested field."""
         to_model_type = self._get_model_type_from_field(to_field)
-        if not to_model_type or not issubclass(to_model_type, BaseModel):
+        if not to_model_type:
             return None
 
         to_info = self.registry.get_model_info(to_model_type)
@@ -337,34 +581,175 @@ class MigrationManager:
             return None
 
         model_name, to_version = to_info
+        from_version = self._get_simple_source_version(from_field, model_name)
+
+        return (model_name, from_version or to_version, to_version)
+
+    def _get_simple_source_version(
+        self, from_field: FieldInfo | None, model_name: str
+    ) -> ModelVersion | None:
+        """Get the source version for a simple nested field."""
+        if not from_field:
+            return None
+
+        from_model_type = self._get_model_type_from_field(from_field)
+        if not from_model_type:
+            return None
+
+        from_info = self.registry.get_model_info(from_model_type)
+        if not from_info or from_info[0] != model_name:
+            return None
+
+        return from_info[1]
+
+    def _get_discriminator_key(self, field: FieldInfo) -> str | None:
+        """Extract the discriminator key from a field."""
+        discriminator = field.discriminator
+        if discriminator is None:
+            return None
+
+        if isinstance(discriminator, str):
+            return discriminator
+
+        if hasattr(discriminator, "discriminator") and isinstance(
+            discriminator.discriminator, str
+        ):
+            return discriminator.discriminator
+
+        return None
+
+    def _find_discriminator_value(
+        self, value: ModelData, discriminator_key: str, field: FieldInfo
+    ) -> Any:
+        """Find the discriminator value in data, checking field name and aliases."""
+        if discriminator_key in value:
+            return value[discriminator_key]
+
+        for model_type in self._get_union_members(field):
+            if discriminator_key not in model_type.model_fields:
+                continue
+
+            disc_field = model_type.model_fields[discriminator_key]
+
+            for alias_attr in ["alias", "serialization_alias"]:
+                alias = getattr(disc_field, alias_attr, None)
+                if alias and alias in value:
+                    return value[alias]
 
-        # Get the source version
-        if from_field:
-            from_model_type = self._get_model_type_from_field(from_field)
-            if from_model_type and issubclass(from_model_type, BaseModel):
-                from_info = self.registry.get_model_info(from_model_type)
-                if from_info and from_info[0] == model_name:
-                    from_version = from_info[1]
-                    return (model_name, from_version, to_version)
+            val_alias = disc_field.validation_alias
+            if isinstance(val_alias, str) and val_alias in value:
+                return value[val_alias]
 
-        # If we can't determine the source version, assume it's the same as target
-        return (model_name, to_version, to_version)
+        return None
+
+    def _get_discriminated_source_version(
+        self,
+        from_field: FieldInfo | None,
+        discriminator_key: str,
+        discriminator_value: Any,
+        model_name: str,
+        default_version: ModelVersion,
+    ) -> ModelVersion:
+        """Get the source version for a discriminated union field."""
+        if not from_field:
+            return default_version
+
+        from_model_type = self._find_discriminated_type(
+            from_field, discriminator_key, discriminator_value
+        )
+        if not from_model_type:
+            return default_version
+
+        from_info = self.registry.get_model_info(from_model_type)
+        if not from_info or from_info[0] != model_name:
+            return default_version
+
+        return from_info[1]
 
-    def _get_model_type_from_field(
-        self: Self, field: FieldInfo
+    def _find_discriminated_type(
+        self, field: FieldInfo, discriminator_key: str, discriminator_value: Any
     ) -> type[BaseModel] | None:
-        """Extract the Pydantic model type from a field.
+        """Find the right type in a discriminated union based on discriminator value."""
+        for model_type in self._get_union_members(field):
+            if self._model_matches_discriminator(
+                model_type, discriminator_key, discriminator_value
+            ):
+                return model_type
+        return None
 
-        Handles Optional, List, and other generic types.
+    def _model_matches_discriminator(
+        self,
+        model_type: type[BaseModel],
+        discriminator_key: str,
+        discriminator_value: Any,
+    ) -> bool:
+        """Check if a model type matches a discriminator value."""
+        if discriminator_key not in model_type.model_fields:
+            return False
 
-        Args:
-            field: The field info to extract from.
+        field_info = model_type.model_fields[discriminator_key]
 
-        Returns:
-            The model type if found, None otherwise.
-        """
+        if self._literal_matches_value(field_info.annotation, discriminator_value):
+            return True
+
+        return (
+            field_info.default is not PydanticUndefined
+            and field_info.default == discriminator_value
+        )
+
+    def _literal_matches_value(self, annotation: Any, value: Any) -> bool:
+        """Check if a Literal type annotation contains a specific value."""
+        if annotation is None:
+            return False
+
+        origin = get_origin(annotation)
+        if origin is not Literal:
+            return False
+
+        literal_values = get_args(annotation)
+        return value in literal_values
+
+    def _get_union_members(self, field: FieldInfo) -> list[type[BaseModel]]:
+        """Extract all BaseModel types from a union field."""
         annotation = field.annotation
+        if annotation is None:
+            return []
 
+        origin = get_origin(annotation)
+
+        if origin is None:
+            if isinstance(annotation, type) and issubclass(annotation, BaseModel):
+                return [annotation]
+            return []
+
+        if not self._is_union_type(origin):
+            return []
+
+        args = get_args(annotation)
+        return [
+            arg
+            for arg in args
+            if arg is not type(None)
+            and isinstance(arg, type)
+            and issubclass(arg, BaseModel)
+        ]
+
+    def _is_union_type(self, origin: Any) -> bool:
+        """Check if an origin type represents a Union."""
+        if origin is Union:
+            return True
+
+        if hasattr(types, "UnionType"):
+            try:
+                return origin is types.UnionType
+            except (ImportError, AttributeError):
+                pass
+
+        return False
+
+    def _get_model_type_from_field(self, field: FieldInfo) -> type[BaseModel] | None:
+        """Extract the Pydantic model type from a field."""
+        annotation = field.annotation
         if annotation is None:
             return None
 

pyrmute/_version.py

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

pyrmute/migration_testing.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Iterator
 from dataclasses import dataclass
-from typing import Self
+from typing import Self, TypeAlias
 
 from .types import ModelData
 
@@ -159,3 +159,6 @@ class MigrationTestResults:
             f"✗ {len(self.failures)} of {total_count} test(s) failed "
             f"({passed_count} passed)"
         )
+
+
+MigrationTestCases: TypeAlias = list[tuple[ModelData, ModelData] | MigrationTestCase]

pyrmute/model_manager.py

@@ -13,6 +13,7 @@ from ._schema_manager import SchemaManager
 from .exceptions import MigrationError, ModelNotFoundError
 from .migration_testing import (
     MigrationTestCase,
+    MigrationTestCases,
     MigrationTestResult,
     MigrationTestResults,
 )
@@ -635,7 +636,7 @@ class ModelManager:
         name: str,
         from_version: str | ModelVersion,
         to_version: str | ModelVersion,
-        test_cases: list[tuple[ModelData, ModelData] | MigrationTestCase],
+        test_cases: MigrationTestCases,
     ) -> MigrationTestResults:
         """Test a migration with multiple test cases.
 

{pyrmute-0.3.0.dist-info → pyrmute-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyrmute
-Version: 0.3.0
+Version: 0.4.0
 Summary: Pydantic model migrations and schemas
 Author-email: Matt Ferrera <mattferrera@gmail.com>
 License: MIT
@@ -56,6 +56,28 @@ through multiple versions.
   support for large datasets
 - **Only one dependency** - Pydantic
 
+## When to Use pyrmute
+
+pyrmute excels at handling schema evolution in production systems:
+
+- **Configuration files** - Upgrade user config files as your CLI/desktop app
+  evolves (`.apprc`, `config.json`, `settings.yaml`)
+- **Message queues & event streams** - Handle messages from multiple service
+  versions publishing different schemas (Kafka, RabbitMQ, SQS)
+- **ETL & data imports** - Import CSV/JSON/Excel files exported over years
+  with evolving structures
+- **ML model serving** - Manage feature schema evolution across model versions
+  and A/B tests
+- **API versioning** - Support multiple API versions with automatic
+  request/response migration
+- **Database migrations** - Transparently migrate legacy data on read without
+  downtime
+- **Data archival** - Process historical data dumps with various schema
+  versions
+
+See the [examples/](examples/) directory for complete, runnable code
+demonstrating these patterns.
+
 ## Help
 
 See [documentation](https://mferrera.github.io/pyrmute/) for complete guides
@@ -232,6 +254,71 @@ results = manager.test_migration(
 assert results.all_passed, f"Migration failed: {results.failures}"
 ```
 
+### Bidirectional Migrations
+
+```python
+# Support both upgrades and downgrades
+@manager.migration("Config", "2.0.0", "1.0.0")
+def downgrade_config(data: ModelData) -> ModelData:
+    """Rollback to v1 format."""
+    return {k: v for k, v in data.items() if k in ["setting1", "setting2"]}
+
+# Useful for:
+# - Rolling back deployments
+# - Normalizing outputs from multiple model versions
+# - Supporting legacy systems during transitions
+```
+
+### Nested Model Migrations
+
+```python
+# Automatically migrates nested Pydantic models
+@manager.model("Address", "1.0.0")
+class AddressV1(BaseModel):
+    street: str
+    city: str
+
+@manager.model("Address", "2.0.0")
+class AddressV2(BaseModel):
+    street: str
+    city: str
+    postal_code: str
+
+@manager.model("User", "2.0.0")
+class UserV2(BaseModel):
+    name: str
+    address: AddressV2  # Nested model
+
+# When migrating User, Address is automatically migrated too
+@manager.migration("Address", "1.0.0", "2.0.0")
+def add_postal_code(data: ModelData) -> ModelData:
+    return {**data, "postal_code": "00000"}
+```
+
+### Discriminated Unions
+
+```python
+from typing import Literal, Union
+from pydantic import Field
+
+# Handle complex type hierarchies
+@manager.model("CreditCard", "1.0.0")
+class CreditCardV1(BaseModel):
+    type: Literal["credit_card"] = "credit_card"
+    card_number: str
+
+@manager.model("PayPal", "1.0.0")
+class PayPalV1(BaseModel):
+    type: Literal["paypal"] = "paypal"
+    email: str
+
+@manager.model("Payment", "1.0.0")
+class PaymentV1(BaseModel):
+    method: Union[CreditCardV1, PayPalV1] = Field(discriminator="type")
+
+# Migrations respect discriminated unions
+```
+
 ### Export JSON Schemas
 
 ```python
@@ -267,79 +354,146 @@ config = manager.migrate({"timeout": 60}, "Config", "1.0.0", "2.0.0")
 # ConfigV2(timeout=60, retries=3)
 ```
 
-## Real-World Example
+## Real-World Examples
+
+### Configuration File Evolution
 
 ```python
-from datetime import datetime
-from pydantic import BaseModel, EmailStr
-from pyrmute import ModelManager, ModelData
+# Your CLI tool evolves over time
+@manager.model("AppConfig", "1.0.0")
+class AppConfigV1(BaseModel):
+    api_key: str
+    debug: bool = False
+
+@manager.model("AppConfig", "2.0.0")
+class AppConfigV2(BaseModel):
+    api_key: str
+    api_endpoint: str = "https://api.example.com"
+    log_level: Literal["DEBUG", "INFO", "ERROR"] = "INFO"
+
+@manager.migration("AppConfig", "1.0.0", "2.0.0")
+def upgrade_config(data: dict) -> dict:
+    return {
+        "api_key": data["api_key"],
+        "api_endpoint": "https://api.example.com",
+        "log_level": "DEBUG" if data.get("debug") else "INFO",
+    }
 
-manager = ModelManager()
+# Load and auto-upgrade user's config file
+def load_config(config_path: Path) -> AppConfigV2:
+    with open(config_path) as f:
+        data = json.load(f)
 
+    version = data.get("_version", "1.0.0")
 
-# API v1: Basic order
-@manager.model("Order", "1.0.0")
-class OrderV1(BaseModel):
-    id: str
-    items: list[str]
-    total: float
-
-
-# API v2: Add customer info
-@manager.model("Order", "2.0.0")
-class OrderV2(BaseModel):
-    id: str
-    items: list[str]
-    total: float
-    customer_email: EmailStr
-
-
-# API v3: Structured items and timestamps
-@manager.model("Order", "3.0.0")
-class OrderItemV3(BaseModel):
-    product_id: str
-    quantity: int
-    price: float
-
-
-@manager.model("Order", "3.0.0")
-class OrderV3(BaseModel):
-    id: str
-    items: list[OrderItemV3]
-    total: float
-    customer_email: EmailStr
-    created_at: datetime
-
-
-# Define migrations
-@manager.migration("Order", "1.0.0", "2.0.0")
-def add_customer_email(data: ModelData) -> ModelData:
-    return {**data, "customer_email": "customer@example.com"}
-
-
-@manager.migration("Order", "2.0.0", "3.0.0")
-def structure_items(data: ModelData) -> ModelData:
-    # Convert simple strings to structured items
-    structured_items = [
-        {
-            "product_id": item,
-            "quantity": 1,
-            "price": data["total"] / len(data["items"])
-        }
-        for item in data["items"]
-    ]
-    return {
-        **data,
-        "items": structured_items,
-        "created_at": datetime.now().isoformat()
-    }
+    # Migrate to current version
+    config = manager.migrate(
+        data,
+        "AppConfig",
+        from_version=version,
+        to_version="2.0.0"
+    )
 
-# Migrate old orders from your database
-old_order = {"id": "123", "items": ["widget", "gadget"], "total": 29.99}
-new_order = manager.migrate(old_order, "Order", "1.0.0", "3.0.0")
-database.save(new_order)
+    # Save upgraded config with version tag
+    with open(config_path, "w") as f:
+        json.dump({**config.model_dump(), "_version": "2.0.0"}, f, indent=2)
+
+    return config
 ```
 
+### Message Queue Consumer
+
+```python
+# Handle messages from multiple service versions
+@manager.model("OrderEvent", "1.0.0")
+class OrderEventV1(BaseModel):
+    order_id: str
+    customer_email: str
+    items: list[dict]  # Unstructured
+
+@manager.model("OrderEvent", "2.0.0")
+class OrderEventV2(BaseModel):
+    order_id: str
+    customer_email: str
+    items: list[OrderItem]  # Structured
+    total: Decimal
+
+def process_message(message: dict, schema_version: str) -> None:
+    # Migrate to current schema regardless of source version
+    event = manager.migrate(
+        message,
+        "OrderEvent",
+        from_version=schema_version,
+        to_version="2.0.0"
+    )
+    # Process with current schema only
+    fulfill_order(event)
+```
+
+### ETL Data Import
+
+```python
+# Import historical exports with evolving schemas
+import csv
+
+def import_customers(file_path: Path, file_version: str) -> None:
+    with open(file_path) as f:
+        reader = csv.DictReader(f)
+
+        # Stream migration for memory efficiency
+        for customer in manager.migrate_batch_streaming(
+            reader,
+            "Customer",
+            from_version=file_version,
+            to_version="3.0.0",
+            chunk_size=1000
+        ):
+            database.save(customer)
+
+# Handle files from different years
+import_customers("exports/2022_customers.csv", "1.0.0")
+import_customers("exports/2023_customers.csv", "2.0.0")
+import_customers("exports/2024_customers.csv", "3.0.0")
+```
+
+### ML Model Serving
+
+```python
+# Route requests to appropriate model versions
+class InferenceService:
+    def predict(self, features: dict, request_version: str) -> BaseModel:
+        # Determine target model version (A/B testing, gradual rollout, etc.)
+        model_version = self.get_model_version(features["user_id"])
+
+        # Migrate request to model's expected format
+        model_input = manager.migrate(
+            features,
+            "PredictionRequest",
+            from_version=request_version,
+            to_version=model_version
+        )
+
+        # Run inference
+        prediction = self.models[model_version].predict(model_input)
+
+        # Normalize output for logging/analytics
+        return manager.migrate(
+            prediction,
+            "PredictionResponse",
+            from_version=model_version,
+            to_version="3.0.0"
+        )
+```
+
+See [examples/](examples/) for complete runnable code:
+- `config_file_migration.py` - CLI/desktop app config file evolution
+- `message_queue_consumer.py` - Kafka/RabbitMQ/SQS consumer handling multiple
+  schemas
+- `etl_data_import.py` - CSV/JSON/Excel import pipeline with historical data
+- `ml_inference_pipeline.py` - ML model serving with feature evolution
+- `advanced_features.py` - Complex Pydantic features (unions, nested models,
+  validators)
+
 ## Contributing
 
 For guidance on setting up a development environment and how to make a

pyrmute-0.4.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+pyrmute/__init__.py,sha256=j2pbMYswL0xR8FZwZDg-qAw6HwFpA3KPhY06_2yc_U0,1132
+pyrmute/_migration_manager.py,sha256=TFws66RsEdKLpvjDDQB1pKgeeyUe5WoutacTaeDsZoE,26154
+pyrmute/_registry.py,sha256=iUjMPd6CYgyvWT8PxZqHWBZnsHrX25fOPDi_-k_QDJs,6124
+pyrmute/_schema_manager.py,sha256=eun8PTL9Gv1XAMVKmE3tYmjdrcf701-IapUXjb6WDL0,12122
+pyrmute/_version.py,sha256=2_0GUP7yBCXRus-qiJKxQD62z172WSs1sQ6DVpPsbmM,704
+pyrmute/exceptions.py,sha256=Q57cUuzzMdkIl5Q0_VyLobpdB0WcrE0ggfC-LBoX2Uo,1681
+pyrmute/migration_testing.py,sha256=fpKT2u7pgPRpswb4PUvbd-fQ3W76svNWvEVYVDmb3Dg,5066
+pyrmute/model_diff.py,sha256=vMa2NTYFqt9E7UYDZH4PQmLcoxQw5Sj-nPpUHB_53Ig,9594
+pyrmute/model_manager.py,sha256=a6ecd-lZ3iliP3lqgCAi7xLeFlBh50kBA-m6gLGKRx4,24585
+pyrmute/model_version.py,sha256=ftNDuJlN3S5ZKQK8DKqqwfBDRiz4rGCYn-aJ3n6Zmqk,2025
+pyrmute/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pyrmute/types.py,sha256=56IH8Rl9AmVh_w3V6PbSSEwaPrBSfc4pYrtcxodvlT0,1187
+pyrmute-0.4.0.dist-info/licenses/LICENSE,sha256=otWInySiZeGwhHqQQ7n7nxM5QBSBe2CzeGEmQDZEz8Q,1119
+pyrmute-0.4.0.dist-info/METADATA,sha256=-KbIQN_INu7ZGfyEO65qmSGsWAXiTHQXR_7F5f-enOQ,14480
+pyrmute-0.4.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+pyrmute-0.4.0.dist-info/top_level.txt,sha256=C8QtzqE6yBHkeewSp1QewvsyeHj_VQLYjSa5HLtMiow,8
+pyrmute-0.4.0.dist-info/RECORD,,

pyrmute-0.3.0.dist-info/RECORD

@@ -1,17 +0,0 @@
-pyrmute/__init__.py,sha256=vB0WBe3CukMBDnK0XP4qqehbMM4z_TUQMcTVPCUyt6Q,1082
-pyrmute/_migration_manager.py,sha256=KregnRUKqF1TC9XIpGAHpQvFlnRTEnp2X6Q2sAay8D4,12489
-pyrmute/_registry.py,sha256=iUjMPd6CYgyvWT8PxZqHWBZnsHrX25fOPDi_-k_QDJs,6124
-pyrmute/_schema_manager.py,sha256=eun8PTL9Gv1XAMVKmE3tYmjdrcf701-IapUXjb6WDL0,12122
-pyrmute/_version.py,sha256=5zTqm8rgXsWYBpB2M3Zw_K1D-aV8wP7NsBLrmMKkrAQ,704
-pyrmute/exceptions.py,sha256=Q57cUuzzMdkIl5Q0_VyLobpdB0WcrE0ggfC-LBoX2Uo,1681
-pyrmute/migration_testing.py,sha256=dOR8BDzmz4mFAI4hFtDUCEMS8Qc8qqD_iOV0qRai-qM,4967
-pyrmute/model_diff.py,sha256=vMa2NTYFqt9E7UYDZH4PQmLcoxQw5Sj-nPpUHB_53Ig,9594
-pyrmute/model_manager.py,sha256=e3UKFo79pkseCUFXIzW2_onu3GYjAnY1FR4JR_QF-Gc,24596
-pyrmute/model_version.py,sha256=ftNDuJlN3S5ZKQK8DKqqwfBDRiz4rGCYn-aJ3n6Zmqk,2025
-pyrmute/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pyrmute/types.py,sha256=56IH8Rl9AmVh_w3V6PbSSEwaPrBSfc4pYrtcxodvlT0,1187
-pyrmute-0.3.0.dist-info/licenses/LICENSE,sha256=otWInySiZeGwhHqQQ7n7nxM5QBSBe2CzeGEmQDZEz8Q,1119
-pyrmute-0.3.0.dist-info/METADATA,sha256=jnRgO76ovFaktYKqq5BMf4tpUc_DYO_drET7c1hPVk0,9580
-pyrmute-0.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-pyrmute-0.3.0.dist-info/top_level.txt,sha256=C8QtzqE6yBHkeewSp1QewvsyeHj_VQLYjSa5HLtMiow,8
-pyrmute-0.3.0.dist-info/RECORD,,