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

pyrmute/__init__.py

@@ -19,8 +19,10 @@ from .migration_testing import (
 from .model_diff import ModelDiff
 from .model_manager import ModelManager
 from .model_version import ModelVersion
+from .schema_config import SchemaConfig
 from .types import (
     JsonSchema,
+    JsonSchemaMode,
     MigrationFunc,
     ModelData,
     NestedModelInfo,
@@ -29,6 +31,7 @@ from .types import (
 __all__ = [
     "InvalidVersionError",
     "JsonSchema",
+    "JsonSchemaMode",
     "MigrationError",
     "MigrationFunc",
     "MigrationManager",
@@ -43,6 +46,7 @@ __all__ = [
     "ModelVersion",
     "NestedModelInfo",
     "Registry",
+    "SchemaConfig",
     "SchemaManager",
     "VersionedModelError",
     "__version__",

pyrmute/_registry.py

@@ -10,11 +10,9 @@ from .exceptions import ModelNotFoundError
 from .model_version import ModelVersion
 from .types import (
     DecoratedBaseModel,
-    JsonSchemaGenerator,
     MigrationMap,
     ModelMetadata,
     ModelName,
-    SchemaGenerators,
     VersionedModels,
 )
 
@@ -28,7 +26,6 @@ class Registry:
     Attributes:
         _models: Dictionary mapping model names to version-model mappings.
         _migrations: Dictionary storing migration functions between versions.
-        _schema_generators: Dictionary storing custom schema generators.
         _model_metadata: Dictionary mapping model classes to (name, version).
         _ref_enabled: Dictionary tracking which models have enable_ref=True.
     """
@@ -37,7 +34,6 @@ class Registry:
         """Initialize the model registry."""
         self._models: dict[ModelName, VersionedModels] = defaultdict(dict)
         self._migrations: dict[ModelName, MigrationMap] = defaultdict(dict)
-        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]] = (
@@ -48,7 +44,6 @@ class Registry:
         self: Self,
         name: ModelName,
         version: str | ModelVersion,
-        schema_generator: JsonSchemaGenerator | None = None,
         enable_ref: bool = False,
         backward_compatible: bool = False,
     ) -> Callable[[type[DecoratedBaseModel]], type[DecoratedBaseModel]]:
@@ -57,7 +52,6 @@ class Registry:
         Args:
             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.
             backward_compatible: If True, this model does not need a migration function
@@ -78,8 +72,6 @@ class Registry:
         def decorator(cls: type[DecoratedBaseModel]) -> type[DecoratedBaseModel]:
             self._models[name][ver] = cls
             self._model_metadata[cls] = (name, ver)
-            if schema_generator:
-                self._schema_generators[name][ver] = schema_generator
             if enable_ref:
                 self._ref_enabled[name].add(ver)
             if backward_compatible:

pyrmute/_schema_manager.py

@@ -1,22 +1,27 @@
-"""Schema manager."""
+"""Schema manager with customizable generation and transformers."""
 
 import json
+from collections import defaultdict
 from pathlib import Path
-from typing import Any, Self, get_args, get_origin
+from typing import Any, Self, cast, get_args, get_origin
 
 from pydantic import BaseModel
 from pydantic.fields import FieldInfo
+from pydantic.json_schema import GenerateJsonSchema
 
 from ._registry import Registry
 from .exceptions import ModelNotFoundError
 from .model_version import ModelVersion
+from .schema_config import SchemaConfig
 from .types import (
     JsonSchema,
     JsonSchemaDefinitions,
+    JsonSchemaGenerator,
     JsonValue,
     ModelMetadata,
     ModelName,
     NestedModelInfo,
+    SchemaTransformer,
 )
 
 
@@ -24,53 +29,217 @@ 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.
+    generators, global configuration, per-call overrides, and schema transformers.
 
     Attributes:
         registry: Reference to the Registry.
+        default_config: Default schema generation configuration.
     """
 
-    def __init__(self: Self, registry: Registry) -> None:
+    def __init__(
+        self: Self, registry: Registry, default_config: SchemaConfig | None = None
+    ) -> None:
         """Initialize the schema manager.
 
         Args:
             registry: Registry instance to use.
+            default_config: Default configuration for schema generation.
         """
         self.registry = registry
+        self.default_config = default_config or SchemaConfig()
+        self._transformers: dict[
+            tuple[ModelName, ModelVersion], list[SchemaTransformer]
+        ] = defaultdict(list)
+
+    def set_default_schema_generator(
+        self: Self, generator: JsonSchemaGenerator | type[GenerateJsonSchema]
+    ) -> None:
+        """Set the default schema generator for all schemas.
+
+        Args:
+            generator: Custom schema generator - either a callable or GenerateJsonSchema
+                class.
+
+        Example (Callable):
+            >>> def custom_gen(model: type[BaseModel]) -> JsonSchema:
+            ...     schema = model.model_json_schema()
+            ...     schema["x-custom"] = True
+            ...     return schema
+            >>>
+            >>> manager.set_default_schema_generator(custom_gen)
+
+        Example (Class):
+            >>> from pydantic.json_schema import GenerateJsonSchema
+            >>>
+            >>> class CustomGenerator(GenerateJsonSchema):
+            ...     def generate(
+            ...         self,
+            ...         schema: Mapping[str, Any],
+            ...         mode: JsonSchemaMode = "validation"
+            ...     ) -> JsonSchema:
+            ...         json_schema = super().generate(schema, mode=mode)
+            ...         json_schema["x-custom"] = True
+            ...         return json_schema
+            >>>
+            >>> manager.set_default_schema_generator(CustomGenerator)
+        """
+        self.default_config.schema_generator = generator
+
+    def register_transformer(
+        self: Self,
+        name: ModelName,
+        version: str | ModelVersion,
+        transformer: SchemaTransformer,
+    ) -> None:
+        """Register a schema transformer for a specific model version.
+
+        Transformers are applied after schema generation, allowing simple
+        post-processing of schemas without needing to customize the generation process
+        itself.
+
+        Args:
+            name: Name of the model.
+            version: Model version.
+            transformer: Function that takes and returns a JsonSchema.
+
+        Example:
+            >>> def add_examples(schema: JsonSchema) -> JsonSchema:
+            ...     schema["examples"] = [{"name": "John", "age": 30}]
+            ...     return schema
+            >>>
+            >>> manager.register_transformer("User", "1.0.0", add_examples)
+        """
+        ver = ModelVersion.parse(version) if isinstance(version, str) else version
+        key = (name, ver)
+        self._transformers[key].append(transformer)
 
     def get_schema(
         self: Self,
         name: ModelName,
         version: str | ModelVersion,
+        config: SchemaConfig | None = None,
+        apply_transformers: bool = True,
         **schema_kwargs: Any,
     ) -> JsonSchema:
         """Get JSON schema for a specific model version.
 
+        Execution order:
+        1. Generate base schema using Pydantic
+        2. Apply custom generator (if configured)
+        3. Apply registered transformers (if any)
+
         Args:
             name: Name of the model.
             version: Semantic version.
-            **schema_kwargs: Additional arguments for schema generation.
+            config: Optional schema configuration (overrides defaults).
+            apply_transformers: If False, skip transformer application.
+            **schema_kwargs: Additional arguments for schema generation (overrides
+                config).
 
         Returns:
             JSON schema dictionary.
+
+        Example:
+            >>> # Use default config
+            >>> schema = manager.get_schema("User", "1.0.0")
+            >>>
+            >>> # Override with custom config
+            >>> config = SchemaConfig(mode="serialization", by_alias=False)
+            >>> schema = manager.get_schema("User", "1.0.0", config=config)
+            >>>
+            >>> # Quick override with kwargs
+            >>> schema = manager.get_schema("User", "1.0.0", mode="serialization")
+            >>>
+            >>> # Get base schema without transformers
+            >>> base_schema = manager.get_schema(
+            ...     "User", "1.0.0",
+            ...     apply_transformers=False
+            ... )
         """
         ver = ModelVersion.parse(version) if isinstance(version, str) else version
         model = self.registry.get_model(name, ver)
 
-        if (
-            name in self.registry._schema_generators
-            and ver in self.registry._schema_generators[name]
-        ):
-            generator = self.registry._schema_generators[name][ver]
-            return generator(model)
+        # Always use the config-based approach
+        final_config = self.default_config
+        if config is not None:
+            final_config = final_config.merge_with(config)
+
+        if schema_kwargs:
+            kwargs_config = SchemaConfig(extra_kwargs=schema_kwargs)
+            final_config = final_config.merge_with(kwargs_config)
+
+        schema: JsonSchema
+        if final_config.is_callable_generator():
+            schema = final_config.schema_generator(model)  # type: ignore
+        else:
+            schema = model.model_json_schema(**final_config.to_kwargs())
+
+        if apply_transformers:
+            key = (name, ver)
+            if key in self._transformers:
+                for transformer in self._transformers[key]:
+                    schema = transformer(schema)
+
+        return schema
+
+    def get_transformers(
+        self: Self,
+        name: ModelName,
+        version: str | ModelVersion,
+    ) -> list[SchemaTransformer]:
+        """Get all transformers registered for a model version.
+
+        Args:
+            name: Name of the model.
+            version: Model version.
+
+        Returns:
+            List of transformer functions.
+        """
+        ver = ModelVersion.parse(version) if isinstance(version, str) else version
+        key = (name, ver)
+        return self._transformers.get(key, [])
+
+    def clear_transformers(
+        self: Self,
+        name: ModelName | None = None,
+        version: str | ModelVersion | None = None,
+    ) -> None:
+        """Clear registered transformers.
 
-        return model.model_json_schema(**schema_kwargs)
+        Args:
+            name: Optional model name. If None, clears all transformers.
+            version: Optional version. If None (but name provided), clears all versions
+                of that model.
+
+        Example:
+            >>> # Clear all transformers
+            >>> manager.clear_transformers()
+            >>>
+            >>> # Clear all User transformers
+            >>> manager.clear_transformers("User")
+            >>>
+            >>> # Clear specific version
+            >>> manager.clear_transformers("User", "1.0.0")
+        """
+        if name is None:
+            self._transformers.clear()
+        elif version is None:
+            keys_to_remove = [key for key in self._transformers if key[0] == name]
+            for key in keys_to_remove:
+                del self._transformers[key]
+        else:
+            ver = ModelVersion.parse(version) if isinstance(version, str) else version
+            key = (name, ver)
+            if key in self._transformers:
+                del self._transformers[key]
 
     def get_schema_with_separate_defs(
         self: Self,
         name: ModelName,
         version: str | ModelVersion,
         ref_template: str = "{model}_v{version}.json",
+        config: SchemaConfig | None = None,
         **schema_kwargs: Any,
     ) -> JsonSchema:
         """Get JSON schema with separate definition files for nested models.
@@ -83,6 +252,7 @@ class SchemaManager:
             version: Semantic version.
             ref_template: Template for generating $ref URLs. Supports {model} and
                 {version} placeholders.
+            config: Optional schema configuration.
             **schema_kwargs: Additional arguments for schema generation.
 
         Returns:
@@ -91,16 +261,16 @@ class SchemaManager:
         Example:
             >>> schema = manager.get_schema_with_separate_defs(
             ...     "User", "2.0.0",
-            ...     ref_template="https://example.com/schemas/{model}_v{version}.json"
+            ...     ref_template="https://example.com/schemas/{model}_v{version}.json",
+            ...     mode="serialization"
             ... )
         """
         ver = ModelVersion.parse(version) if isinstance(version, str) else version
-        schema = self.get_schema(name, ver, **schema_kwargs)
+        schema = self.get_schema(name, ver, config=config, **schema_kwargs)
 
-        # Extract and replace definitions with external references
         if "$defs" in schema or "definitions" in schema:
             defs_key = "$defs" if "$defs" in schema else "definitions"
-            definitions: JsonSchemaDefinitions = schema.pop(defs_key, {})  # type: ignore[assignment]
+            definitions = cast("JsonSchemaDefinitions", schema.pop(defs_key, {}))
 
             # Update all $ref in the schema to point to external files
             schema = self._replace_refs_with_external(schema, definitions, ref_template)
@@ -209,11 +379,14 @@ class SchemaManager:
                     return (name, version)
         return None
 
-    def get_all_schemas(self: Self, name: ModelName) -> dict[ModelVersion, JsonSchema]:
+    def get_all_schemas(
+        self: Self, name: ModelName, config: SchemaConfig | None = None
+    ) -> dict[ModelVersion, JsonSchema]:
         """Get all schemas for a model across all versions.
 
         Args:
             name: Name of the model.
+            config: Optional schema configuration.
 
         Returns:
             Dictionary mapping versions to their schemas.
@@ -225,7 +398,7 @@ class SchemaManager:
             raise ModelNotFoundError(name)
 
         return {
-            version: self.get_schema(name, version)
+            version: self.get_schema(name, version, config=config)
             for version in self.registry._models[name]
         }
 
@@ -235,6 +408,7 @@ class SchemaManager:
         indent: int = 2,
         separate_definitions: bool = False,
         ref_template: str | None = None,
+        config: SchemaConfig | None = None,
     ) -> None:
         """Dump all schemas to JSON files.
 
@@ -245,27 +419,24 @@ class SchemaManager:
                 models that have enable_ref=True.
             ref_template: Template for $ref URLs when separate_definitions=True.
                 Defaults to relative file references if not provided.
+            config: Optional schema configuration for all exported schemas.
 
         Example:
-            >>> # Inline definitions (default)
-            >>> manager.dump_schemas("schemas/")
-            >>>
-            >>> # Separate sub-schemas with relative refs (when enable_ref=True models)
-            >>> manager.dump_schemas("schemas/", separate_definitions=True)
-            >>>
-            >>> # Separate sub-schemas with absolute URLs
-            >>> manager.dump_schemas(
-            ...     "schemas/",
-            ...     separate_definitions=True,
-            ...     ref_template="https://example.com/schemas/{model}_v{version}.json"
+            >>> # Export with custom schema generator
+            >>> config = SchemaConfig(
+            ...     schema_generator=CustomGenerator,
+            ...     mode="serialization"
             ... )
+            >>> manager.dump_schemas("schemas/", config=config)
         """
         output_path = Path(output_dir)
         output_path.mkdir(parents=True, exist_ok=True)
 
         if not separate_definitions:
             for name in self.registry._models:
-                for version, schema in self.get_all_schemas(name).items():
+                for version, schema in self.get_all_schemas(
+                    name, config=config
+                ).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)
@@ -276,7 +447,7 @@ class SchemaManager:
             for name in self.registry._models:
                 for version in self.registry._models[name]:
                     schema = self.get_schema_with_separate_defs(
-                        name, version, ref_template
+                        name, version, ref_template, config=config
                     )
                     file_path = output_path / f"{name}_v{version}.json"
                     with open(file_path, "w", encoding="utf-8") as f:

pyrmute/_version.py

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

pyrmute/model_manager.py

@@ -1,4 +1,4 @@
-"""Model manager."""
+"""ModelManager class."""
 
 from collections.abc import Callable, Iterable
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Any, Self
 
 from pydantic import BaseModel
+from pydantic.json_schema import GenerateJsonSchema
 
 from ._migration_manager import MigrationManager
 from ._registry import Registry
@@ -19,6 +20,7 @@ from .migration_testing import (
 )
 from .model_diff import ModelDiff
 from .model_version import ModelVersion
+from .schema_config import SchemaConfig
 from .types import (
     DecoratedBaseModel,
     JsonSchema,
@@ -26,20 +28,16 @@ from .types import (
     MigrationFunc,
     ModelData,
     NestedModelInfo,
+    SchemaTransformer,
 )
 
 
 class ModelManager:
-    """High-level interface for versioned model management.
+    """High-level interface for versioned model management and schema generation.
 
     ModelManager provides a unified API for managing schema evolution across different
     versions of Pydantic models. It handles model registration, automatic migration
-    between versions, schema generation, and batch processing operations.
-
-    Attributes:
-        registry: Registry instance managing all registered model versions.
-        migration_manager: MigrationManager instance handling migration logic and paths.
-        schema_manager: SchemaManager instance for JSON schema generation and export.
+    between versions, customizable schema generation, and batch processing operations.
 
     Basic Usage:
         >>> manager = ModelManager()
@@ -64,6 +62,56 @@ class ModelManager:
         >>> user = manager.migrate(old_data, "User", "1.0.0", "2.0.0")
         >>> # Result: UserV2(name="Alice", email="unknown@example.com")
 
+    Custom Schema Generation:
+        >>> from pydantic.json_schema import GenerateJsonSchema
+        >>>
+        >>> class CustomSchemaGenerator(GenerateJsonSchema):
+        ...     '''Add custom metadata to all schemas.'''
+        ...     def generate(
+        ...         self,
+        ...         schema: Mapping[str, Any],
+        ...         mode: JsonSchemaMode = "validation"
+        ...     ) -> JsonSchema:
+        ...         json_schema = super().generate(schema, mode=mode)
+        ...         json_schema["x-company"] = "Acme"
+        ...         json_schema["$schema"] = self.schema_dialect
+        ...         return json_schema
+        >>>
+        >>> # Set at manager level (applies to all schemas)
+        >>> manager = ModelManager(
+        ...     default_schema_config=SchemaConfig(
+        ...         schema_generator=CustomSchemaGenerator,
+        ...         mode="validation",
+        ...         by_alias=True
+        ...     )
+        ... )
+        >>>
+        >>> @manager.model("User", "1.0.0")
+        ... class User(BaseModel):
+        ...     name: str = Field(title="Full Name")
+        ...     email: str
+        >>>
+        >>> # Get schema with default config
+        >>> schema = manager.get_schema("User", "1.0.0")
+        >>> # Will include x-company: 'Acme'
+
+    Schema Transformers:
+        >>> manager = ModelManager()
+        >>>
+        >>> @manager.model("Product", "1.0.0")
+        ... class Product(BaseModel):
+        ...     name: str
+        ...     price: float
+        >>>
+        >>> # Add transformer for specific model
+        >>> @manager.schema_transformer("Product", "1.0.0")
+        ... def add_examples(schema: JsonSchema) -> JsonSchema:
+        ...     schema["examples"] = [{"name": "Widget", "price": 9.99}]
+        ...     return schema
+        >>>
+        >>> schema = manager.get_schema("Product", "1.0.0")
+        >>> # Will include examples
+
     Advanced Features:
         >>> # Batch migration with parallel processing
         >>> users = manager.migrate_batch(
@@ -72,7 +120,9 @@ class ModelManager:
         ... )
         >>>
         >>> # Stream large datasets efficiently
-        >>> for user in manager.migrate_batch_streaming(large_dataset, "User", "1.0.0", "2.0.0"):
+        >>> for user in manager.migrate_batch_streaming(
+        ...     large_dataset, "User", "1.0.0", "2.0.0"
+        ...  ):
         ...     save_to_database(user)
         >>>
         >>> # Compare versions and export schemas
@@ -84,23 +134,32 @@ class ModelManager:
         >>> results = manager.test_migration(
         ...     "User", "1.0.0", "2.0.0",
         ...     test_cases=[
-        ...         ({"name": "Alice"}, {"name": "Alice", "email": "unknown@example.com"})
+        ...         (
+        ...              {"name": "Alice"},
+        ...              {"name": "Alice", "email": "unknown@example.com"}
+        ...         )
         ...     ]
         ... )
         >>> results.assert_all_passed()
-    """  # noqa: E501
+    """
+
+    def __init__(self: Self, default_schema_config: SchemaConfig | None = None) -> None:
+        """Initialize the versioned model manager.
 
-    def __init__(self: Self) -> None:
-        """Initialize the versioned model manager."""
+        Args:
+            default_schema_config: Default configuration for schema generation
+                applied to all schema operations unless overridden.
+        """
         self._registry = Registry()
         self._migration_manager = MigrationManager(self._registry)
-        self._schema_manager = SchemaManager(self._registry)
+        self._schema_manager = SchemaManager(
+            self._registry, default_config=default_schema_config
+        )
 
     def model(
         self: Self,
         name: str,
         version: str | ModelVersion,
-        schema_generator: JsonSchemaGenerator | None = None,
         enable_ref: bool = False,
         backward_compatible: bool = False,
     ) -> Callable[[type[DecoratedBaseModel]], type[DecoratedBaseModel]]:
@@ -109,7 +168,6 @@ class ModelManager:
         Args:
             name: Name of the model.
             version: Semantic version.
-            schema_generator: Optional custom schema generator.
             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
@@ -130,9 +188,7 @@ class ModelManager:
             ... class CityV1(BaseModel):
             ...     city: City
         """
-        return self._registry.register(
-            name, version, schema_generator, enable_ref, backward_compatible
-        )
+        return self._registry.register(name, version, enable_ref, backward_compatible)
 
     def migration(
         self: Self,
@@ -309,27 +365,12 @@ class ModelManager:
             from_version: Source version.
             to_version: Target version.
             parallel: If True, use parallel processing.
-            max_workers: Maximum number of workers for parallel processing.  Defaults to
-                None (uses executor default).
+            max_workers: Maximum number of workers for parallel processing.
             use_processes: If True, use ProcessPoolExecutor instead of
-                ThreadPoolExecutor. Useful for CPU-intensive migrations.
+                ThreadPoolExecutor.
 
         Returns:
             List of migrated BaseModel instances.
-
-        Example:
-            >>> legacy_users = [
-            ...     {"name": "Alice"},
-            ...     {"name": "Bob"},
-            ...     {"name": "Charlie"}
-            ... ]
-            >>> users = manager.migrate_batch(
-            ...     legacy_users,
-            ...     "User",
-            ...     from_version="1.0.0",
-            ...     to_version="3.0.0",
-            ...     parallel=True
-            ... )
         """
         data_list = list(data_list)
 
@@ -372,15 +413,6 @@ class ModelManager:
 
         Returns:
             List of raw migrated dictionaries.
-
-        Example:
-            >>> legacy_data = [{"name": "Alice"}, {"name": "Bob"}]
-            >>> migrated_data = manager.migrate_batch_data(
-            ...     legacy_data,
-            ...     "User",
-            ...     from_version="1.0.0",
-            ...     to_version="2.0.0"
-            ... )
         """
         data_list = list(data_list)
 
@@ -411,9 +443,6 @@ class ModelManager:
     ) -> Iterable[BaseModel]:
         """Migrate data in chunks, yielding results as they complete.
 
-        Useful for large datasets where you want to start processing results before all
-        migrations complete.
-
         Args:
             data_list: Iterable of data dictionaries to migrate.
             name: Name of the model.
@@ -423,17 +452,6 @@ class ModelManager:
 
         Yields:
             Migrated BaseModel instances.
-
-        Example:
-            >>> legacy_users = load_large_dataset()
-            >>> for user in manager.migrate_batch_streaming(
-            ...     legacy_users,
-            ...     "User",
-            ...     from_version="1.0.0",
-            ...     to_version="3.0.0"
-            ... ):
-            ...     # Process each user as it's migrated
-            ...     save_to_database(user)
         """
         chunk = []
 
@@ -457,9 +475,6 @@ class ModelManager:
     ) -> Iterable[ModelData]:
         """Migrate data in chunks, yielding raw dictionaries as they complete.
 
-        Useful for large datasets where you want to start processing results before all
-        migrations complete, without the validation overhead.
-
         Args:
             data_list: Iterable of data dictionaries to migrate.
             name: Name of the model.
@@ -469,17 +484,6 @@ class ModelManager:
 
         Yields:
             Raw migrated dictionaries.
-
-        Example:
-            >>> legacy_data = load_large_dataset()
-            >>> for data in manager.migrate_batch_data_streaming(
-            ...     legacy_data,
-            ...     "User",
-            ...     from_version="1.0.0",
-            ...     to_version="3.0.0"
-            ... ):
-            ...     # Process raw data as it's migrated
-            ...     bulk_insert_to_database(data)
         """
         chunk = []
 
@@ -503,9 +507,6 @@ class ModelManager:
     ) -> ModelDiff:
         """Get a detailed diff between two model versions.
 
-        Compares field names, types, requirements, and default values to provide a
-        comprehensive view of what changed between versions.
-
         Args:
             name: Name of the model.
             from_version: Source version.
@@ -513,12 +514,6 @@ class ModelManager:
 
         Returns:
             ModelDiff with detailed change information.
-
-        Example:
-            >>> diff = manager.diff("User", "1.0.0", "2.0.0")
-            >>> print(diff.to_markdown())
-            >>> print(f"Added: {diff.added_fields}")
-            >>> print(f"Removed: {diff.removed_fields}")
         """
         from_ver_str = str(
             ModelVersion.parse(from_version)
@@ -542,10 +537,136 @@ class ModelManager:
             to_version=to_ver_str,
         )
 
+    def set_default_schema_generator(
+        self: Self, generator: JsonSchemaGenerator | type[GenerateJsonSchema]
+    ) -> None:
+        """Set the default schema generator for all schemas.
+
+        This is a convenience method that updates the default schema configuration.
+
+        Args:
+            generator: Custom schema generator - either a callable or GenerateJsonSchema
+                class.
+
+        Example (Callable):
+            >>> def my_generator(model: type[BaseModel]) -> JsonSchema:
+            ...     schema = model.model_json_schema()
+            ...     schema["x-custom"] = True
+            ...     return schema
+            >>>
+            >>> manager = ModelManager()
+            >>> manager.set_default_schema_generator(my_generator)
+
+        Example (Class - Recommended):
+            >>> from pydantic.json_schema import GenerateJsonSchema
+            >>>
+            >>> class MyGenerator(GenerateJsonSchema):
+            ...     def generate(
+            ...         self,
+            ...         schema: Mapping[str, Any],
+            ...         mode: JsonSchemaMode = "validation"
+            ...     ) -> JsonSchema:
+            ...         json_schema = super().generate(schema, mode=mode)
+            ...         json_schema["x-custom"] = True
+            ...         json_schema["$schema"] = self.schema_dialect
+            ...         return json_schema
+            >>>
+            >>> manager = ModelManager()
+            >>> manager.set_default_schema_generator(MyGenerator)
+            >>>
+            >>> # All subsequent schema calls will use MyGenerator
+            >>> schema = manager.get_schema("User", "1.0.0")
+        """
+        self._schema_manager.set_default_schema_generator(generator)
+
+    def schema_transformer(
+        self: Self,
+        name: str,
+        version: str | ModelVersion,
+    ) -> Callable[[SchemaTransformer], SchemaTransformer]:
+        """Decorator to register a schema transformer for a model version.
+
+        Transformers are simple functions that modify a schema after generation.
+        They're useful for model-specific customizations that don't require deep
+        integration with Pydantic's generation process.
+
+        Args:
+            name: Name of the model.
+            version: Model version.
+
+        Returns:
+            Decorator function.
+
+        Example:
+            >>> @manager.schema_transformer("User", "1.0.0")
+            ... def add_auth_metadata(schema: JsonSchema) -> JsonSchema:
+            ...     schema["x-requires-auth"] = True
+            ...     schema["x-auth-level"] = 'admin'
+            ...     return schema
+            >>>
+            >>> @manager.schema_transformer("Product", "2.0.0")
+            ... def add_product_examples(schema: JsonSchema) -> JsonSchema:
+            ...     schema["examples"] = [
+            ...         {"name": "Widget", "price": 9.99},
+            ...         {"name": "Gadget", "price": 19.99}
+            ...     ]
+            ...     return schema
+        """
+
+        def decorator(func: SchemaTransformer) -> SchemaTransformer:
+            self._schema_manager.register_transformer(name, version, func)
+            return func
+
+        return decorator
+
+    def get_schema_transformers(
+        self: Self,
+        name: str,
+        version: str | ModelVersion,
+    ) -> list[SchemaTransformer]:
+        """Get all transformers for a model version.
+
+        Args:
+            name: Name of the model.
+            version: Model version.
+
+        Returns:
+            List of transformer functions.
+
+        Example:
+            >>> transformers = manager.get_schema_transformers("User", "1.0.0")
+            >>> print(f"Found {len(transformers)} transformers")
+        """
+        return self._schema_manager.get_transformers(name, version)
+
+    def clear_schema_transformers(
+        self: Self,
+        name: str | None = None,
+        version: str | ModelVersion | None = None,
+    ) -> None:
+        """Clear schema transformers.
+
+        Args:
+            name: Optional model name. If None, clears all.
+            version: Optional version. If None, clears all versions of model.
+
+        Example:
+            >>> # Clear all transformers
+            >>> manager.clear_schema_transformers()
+            >>>
+            >>> # Clear User transformers
+            >>> manager.clear_schema_transformers("User")
+            >>>
+            >>> # Clear specific version
+            >>> manager.clear_schema_transformers("User", "1.0.0")
+        """
+        self._schema_manager.clear_transformers(name, version)
+
     def get_schema(
         self: Self,
         name: str,
         version: str | ModelVersion,
+        config: SchemaConfig | None = None,
         **kwargs: Any,
     ) -> JsonSchema:
         """Get JSON schema for a specific version.
@@ -553,12 +674,25 @@ class ModelManager:
         Args:
             name: Name of the model.
             version: Semantic version.
-            **kwargs: Additional schema generation arguments.
+            config: Optional schema configuration (overrides default).
+            **kwargs: Additional schema generation arguments (e.g.,
+                mode="serialization").
 
         Returns:
             JSON schema dictionary.
+
+        Example:
+            >>> # Use default config
+            >>> schema = manager.get_schema("User", "1.0.0")
+            >>>
+            >>> # Override with custom config
+            >>> config = SchemaConfig(mode="serialization")
+            >>> schema = manager.get_schema("User", "1.0.0", config=config)
+            >>>
+            >>> # Quick override with kwargs
+            >>> schema = manager.get_schema("User", "1.0.0", mode="serialization")
         """
-        return self._schema_manager.get_schema(name, version, **kwargs)
+        return self._schema_manager.get_schema(name, version, config=config, **kwargs)
 
     def list_models(self: Self) -> list[str]:
         """Get list of all registered models.
@@ -585,6 +719,7 @@ class ModelManager:
         indent: int = 2,
         separate_definitions: bool = False,
         ref_template: str | None = None,
+        config: SchemaConfig | None = None,
     ) -> None:
         """Export all schemas to JSON files.
 
@@ -592,27 +727,30 @@ class ModelManager:
             output_dir: Directory path for output.
             indent: JSON indentation level.
             separate_definitions: If True, create separate schema files for nested
-                models and use $ref to reference them. Only applies to models with
-                'enable_ref=True'.
+                models and use $ref to reference them.
             ref_template: Template for $ref URLs when separate_definitions=True.
-                Defaults to relative file references if not provided.
+            config: Optional schema configuration for all exported schemas.
 
         Example:
-            >>> # Inline definitions (default)
-            >>> manager.dump_schemas("schemas/")
-            >>>
-            >>> # Separate sub-schemas with relative refs
-            >>> manager.dump_schemas("schemas/", separate_definitions=True)
+            >>> # Export with custom generator
+            >>> config = SchemaConfig(
+            ...     schema_generator=CustomGenerator,
+            ...     mode="validation"
+            ... )
+            >>> manager.dump_schemas("schemas/", config=config)
             >>>
-            >>> # Separate sub-schemas with absolute URLs
+            >>> # Export validation and serialization schemas separately
             >>> manager.dump_schemas(
-            ...     "schemas/",
-            ...     separate_definitions=True,
-            ...     ref_template="https://example.com/schemas/{model}_v{version}.json"
+            ...     "schemas/validation/",
+            ...     config=SchemaConfig(mode="validation")
+            ... )
+            >>> manager.dump_schemas(
+            ...     "schemas/serialization/",
+            ...     config=SchemaConfig(mode="serialization")
             ... )
         """
         self._schema_manager.dump_schemas(
-            output_dir, indent, separate_definitions, ref_template
+            output_dir, indent, separate_definitions, ref_template, config=config
         )
 
     def get_nested_models(
@@ -640,54 +778,15 @@ class ModelManager:
     ) -> MigrationTestResults:
         """Test a migration with multiple test cases.
 
-        Executes a migration on multiple test inputs and validates the outputs match
-        expected values. Useful for regression testing and validating migration logic.
-
         Args:
             name: Name of the model.
             from_version: Source version to migrate from.
             to_version: Target version to migrate to.
-            test_cases: List of test cases, either as (source, target) tuples or
-                MigrationTestCase objects. If target is None, only verifies the
-                migration completes without errors.
+            test_cases: List of test cases.
 
         Returns:
             MigrationTestResults containing individual results for each test case.
-
-        Example:
-            >>> # Using tuples (source, target)
-            >>> results = manager.test_migration(
-            ...     "User", "1.0.0", "2.0.0",
-            ...     test_cases=[
-            ...         ({"name": "Alice"}, {"name": "Alice", "email": "alice@example.com"}),
-            ...         ({"name": "Bob"}, {"name": "Bob", "email": "bob@example.com"})
-            ...     ]
-            ... )
-            >>> assert results.all_passed
-            >>>
-            >>> # Using MigrationTestCase objects
-            >>> results = manager.test_migration(
-            ...     "User", "1.0.0", "2.0.0",
-            ...     test_cases=[
-            ...         MigrationTestCase(
-            ...             source={"name": "Alice"},
-            ...             target={"name": "Alice", "email": "alice@example.com"},
-            ...             description="Standard user migration"
-            ...         )
-            ...     ]
-            ... )
-            >>>
-            >>> # Use in pytest
-            >>> def test_user_migration():
-            ...     results = manager.test_migration("User", "1.0.0", "2.0.0", test_cases)
-            ...     results.assert_all_passed()  # Raises AssertionError with details if failed
-            >>>
-            >>> # Inspect failures
-            >>> if not results.all_passed:
-            ...     for failure in results.failures:
-            ...         print(f"Failed: {failure.test_case.description}")
-            ...         print(f"  Error: {failure.error}")
-        """  # noqa: E501
+        """
         results = []
 
         for test_case_input in test_cases:

pyrmute/schema_config.py

@@ -0,0 +1,130 @@
+"""Schema configuration for customized schema generation."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Self
+
+from pydantic.json_schema import GenerateJsonSchema
+
+if TYPE_CHECKING:
+    from .types import JsonSchemaGenerator, JsonSchemaMode
+
+
+@dataclass
+class SchemaConfig:
+    """Configuration for JSON schema generation.
+
+    This class provides fine-grained control over how Pydantic generates JSON schemas,
+    supporting both callable generators and Pydantic's GenerateJsonSchema classes.
+
+    Attributes:
+        schema_generator: Custom schema generator. Can be either:
+            - A callable taking (type[BaseModel]) -> JsonSchema
+            - A subclass of pydantic.json_schema.GenerateJsonSchema
+        mode: Schema generation mode - 'validation' for input validation or
+            'serialization' for output serialization.
+        by_alias: Whether to use field aliases in the schema.
+        ref_template: Template for JSON schema $ref URIs.
+        extra_kwargs: Additional arguments to pass to model_json_schema().
+
+    Example (Callable Generator):
+        >>> def custom_generator(model: type[BaseModel]) -> JsonSchema:
+        ...     schema = model.model_json_schema()
+        ...     schema["x-custom"] = "metadata"
+        ...     return schema
+        >>>
+        >>> config = SchemaConfig(
+        ...     schema_generator=custom_generator,
+        ...     mode="validation"
+        ... )
+
+    Example (GenerateJsonSchema Class):
+        >>> from pydantic.json_schema import GenerateJsonSchema
+        >>>
+        >>> class CustomSchemaGenerator(GenerateJsonSchema):
+        ...     def generate(
+        ...         self,
+        ...         schema: Mapping[str, Any],
+        ...         mode: JsonSchemaMode = "validation"
+        ...     ) -> JsonSchema:
+        ...         json_schema = super().generate(schema, mode=mode)
+        ...         json_schema["x-custom"] = "metadata"
+        ...         return json_schema
+        >>>
+        >>> config = SchemaConfig(
+        ...     schema_generator=CustomSchemaGenerator,
+        ...     mode="validation",
+        ...     by_alias=True
+        ... )
+    """
+
+    schema_generator: JsonSchemaGenerator | type[GenerateJsonSchema] | None = None
+    mode: JsonSchemaMode = "validation"
+    by_alias: bool = True
+    ref_template: str = "#/$defs/{model}"
+    extra_kwargs: dict[str, Any] = field(default_factory=dict)
+
+    def merge_with(self: Self, other: SchemaConfig | None) -> SchemaConfig:
+        """Merge this config with another, with other taking precedence.
+
+        Args:
+            other: Configuration to merge with (overrides this config).
+
+        Returns:
+            New SchemaConfig with merged values.
+        """
+        if other is None:
+            return self
+
+        return SchemaConfig(
+            schema_generator=other.schema_generator or self.schema_generator,
+            mode=other.mode if other.mode != "validation" else self.mode,
+            by_alias=other.by_alias if not other.by_alias else self.by_alias,
+            ref_template=other.ref_template
+            if other.ref_template != "#/$defs/{model}"
+            else self.ref_template,
+            extra_kwargs={**self.extra_kwargs, **other.extra_kwargs},
+        )
+
+    def to_kwargs(self: Self) -> dict[str, Any]:
+        """Convert config to kwargs for model_json_schema().
+
+        Note: If schema_generator is a callable (JsonSchemaGenerator type), it cannot be
+        passed to model_json_schema() and must be handled separately by calling it
+        directly.
+
+        Returns:
+            Dictionary of arguments for Pydantic's model_json_schema(). If
+            schema_generator is a callable, it will NOT be included.
+        """
+        kwargs = {
+            "mode": self.mode,
+            "by_alias": self.by_alias,
+            "ref_template": self.ref_template,
+            **self.extra_kwargs,
+        }
+
+        # Only add schema_generator if it's a GenerateJsonSchema class
+        # Callable generators are handled separately
+        if (
+            self.schema_generator is not None
+            and isinstance(self.schema_generator, type)
+            and issubclass(self.schema_generator, GenerateJsonSchema)
+        ):
+            kwargs["schema_generator"] = self.schema_generator
+
+        return kwargs
+
+    def is_callable_generator(self: Self) -> bool:
+        """Check if schema_generator is a callable function.
+
+        Returns:
+            True if schema_generator is a callable (not a class).
+        """
+        if self.schema_generator is None:
+            return False
+
+        return callable(self.schema_generator) and not isinstance(
+            self.schema_generator, type
+        )

pyrmute/types.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, TypeAlias, TypeVar
+from typing import Any, Literal, TypeAlias, TypeVar
 
 from pydantic import BaseModel
 
@@ -16,9 +16,10 @@ JsonValue: TypeAlias = (
     int | float | str | bool | None | list["JsonValue"] | dict[str, "JsonValue"]
 )
 JsonSchema: TypeAlias = dict[str, JsonValue]
+JsonSchemaMode = Literal["validation", "serialization"]
 JsonSchemaDefinitions: TypeAlias = dict[str, JsonValue]
 JsonSchemaGenerator: TypeAlias = Callable[[type[BaseModel]], JsonSchema]
-SchemaGenerators: TypeAlias = dict[ModelVersion, JsonSchemaGenerator]
+SchemaTransformer = Callable[[JsonSchema], JsonSchema]
 
 ModelData: TypeAlias = dict[str, Any]
 MigrationFunc: TypeAlias = Callable[[ModelData], ModelData]

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyrmute
-Version: 0.4.0
+Version: 0.5.0
 Summary: Pydantic model migrations and schemas
 Author-email: Matt Ferrera <mattferrera@gmail.com>
 License: MIT
@@ -58,7 +58,7 @@ through multiple versions.
 
 ## When to Use pyrmute
 
-pyrmute excels at handling schema evolution in production systems:
+pyrmute is useful for handling schema evolution in production systems:
 
 - **Configuration files** - Upgrade user config files as your CLI/desktop app
   evolves (`.apprc`, `config.json`, `settings.yaml`)
@@ -78,6 +78,21 @@ pyrmute excels at handling schema evolution in production systems:
 See the [examples/](examples/) directory for complete, runnable code
 demonstrating these patterns.
 
+## When Not to Use
+
+pyrmute may not be the right choice if you have:
+
+- **High-throughput systems** - Runtime migration adds latency to hot paths.
+  Use upfront batch migrations instead.
+- **Multi-language services** - Python-only. Use Protobuf, Avro, or JSON
+  Schema for polyglot architectures.
+- **Existing schema registries** - Already using Confluent/AWS Glue? Stick
+  with them for compatibility enforcement and governance.
+- **Stable schemas** - Models rarely change? Traditional migration tools are
+  simpler and more maintainable.
+- **Database DDL changes** - pyrmute transforms data, not database schemas.
+  Alembic/Flyway or other ORMs may still be needed to alter tables.
+
 ## Help
 
 See [documentation](https://mferrera.github.io/pyrmute/) for complete guides

pyrmute-0.5.0.dist-info/RECORD

@@ -0,0 +1,18 @@
+pyrmute/__init__.py,sha256=vgq5e3jmr2FHmSnnnf-Kvl4Y6oGra3BDl5p_CaHDYLQ,1234
+pyrmute/_migration_manager.py,sha256=TFws66RsEdKLpvjDDQB1pKgeeyUe5WoutacTaeDsZoE,26154
+pyrmute/_registry.py,sha256=eEhMagcpUeZSPNZlIAuTzXcCkPGHTxfaIfEYEoEiFU8,5680
+pyrmute/_schema_manager.py,sha256=9NmyooG92gcOMOYNLRuC8isGUvivIf_Hp8kdVC8lmu8,18622
+pyrmute/_version.py,sha256=fvHpBU3KZKRinkriKdtAt3crenOyysELF-M9y3ozg3U,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=EC3xN4Jve77Fk3xJjC3KrU_sY3Eoae-ucJE8wQ3bO3M,27778
+pyrmute/model_version.py,sha256=ftNDuJlN3S5ZKQK8DKqqwfBDRiz4rGCYn-aJ3n6Zmqk,2025
+pyrmute/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pyrmute/schema_config.py,sha256=ttM4a1-EUUlxmsol6L4rN_Mse-77Y-yidn0ezojy0uY,4767
+pyrmute/types.py,sha256=rzRxOYfh4WPVR1KoNT3vC3UjuBlTarMnNL6Z1Y5icrw,1237
+pyrmute-0.5.0.dist-info/licenses/LICENSE,sha256=otWInySiZeGwhHqQQ7n7nxM5QBSBe2CzeGEmQDZEz8Q,1119
+pyrmute-0.5.0.dist-info/METADATA,sha256=QOv8NMAZ_sw-k2bFOHXM79XGnPUBArt4PRELq5cUdMY,15169
+pyrmute-0.5.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+pyrmute-0.5.0.dist-info/top_level.txt,sha256=C8QtzqE6yBHkeewSp1QewvsyeHj_VQLYjSa5HLtMiow,8
+pyrmute-0.5.0.dist-info/RECORD,,

pyrmute-0.4.0.dist-info/RECORD

@@ -1,17 +0,0 @@
-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,,