pyrmute 0.1.0__py3-none-any.whl

pyrmute/_schema_manager.py

@@ -0,0 +1,350 @@
+"""Schema manager."""
+
+import json
+from pathlib import Path
+from typing import Any, Self, get_args, get_origin
+
+from pydantic import BaseModel
+from pydantic.fields import FieldInfo
+
+from ._registry import Registry
+from .model_version import ModelVersion
+from .types import (
+    JsonSchema,
+    JsonSchemaDefinitions,
+    JsonValue,
+    ModelMetadata,
+    ModelName,
+)
+
+
+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.
+
+    Attributes:
+        registry: Reference to the Registry.
+    """
+
+    def __init__(self: Self, registry: Registry) -> None:
+        """Initialize the schema manager.
+
+        Args:
+            registry: Registry instance to use.
+        """
+        self.registry = registry
+
+    def get_schema(
+        self: Self,
+        name: ModelName,
+        version: str | ModelVersion,
+        **schema_kwargs: Any,
+    ) -> JsonSchema:
+        """Get JSON schema for a specific model version.
+
+        Args:
+            name: Name of the model.
+            version: Semantic version.
+            **schema_kwargs: Additional arguments for schema generation.
+
+        Returns:
+            JSON schema dictionary.
+        """
+        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)
+
+        return model.model_json_schema(**schema_kwargs)
+
+    def get_schema_with_separate_defs(
+        self: Self,
+        name: ModelName,
+        version: str | ModelVersion,
+        ref_template: str = "{model}_v{version}.json",
+        **schema_kwargs: Any,
+    ) -> JsonSchema:
+        """Get JSON schema with separate definition files for nested models.
+
+        This creates a schema where nested Pydantic models are referenced
+        as external JSON schema files rather than inline definitions.
+
+        Args:
+            name: Name of the model.
+            version: Semantic version.
+            ref_template: Template for generating $ref URLs. Supports {model}
+                and {version} placeholders.
+            **schema_kwargs: Additional arguments for schema generation.
+
+        Returns:
+            JSON schema dictionary with external $ref for nested models.
+
+        Example:
+            >>> schema = manager.get_schema_with_separate_defs(
+            ...     "User", "2.0.0",
+            ...     ref_template="https://example.com/schemas/{model}_v{version}.json"
+            ... )
+        """
+        ver = ModelVersion.parse(version) if isinstance(version, str) else version
+
+        # Get the base schema with definitions
+        schema = self.get_schema(name, ver, **schema_kwargs)
+
+        # Extract and replace definitions with external references
+        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]
+
+            # Update all $ref in the schema to point to external files
+            schema = self._replace_refs_with_external(schema, definitions, ref_template)
+
+            # Re-add definitions that weren't converted to external refs
+            remaining_defs = self._get_remaining_defs(schema, definitions)
+            if remaining_defs:
+                schema[defs_key] = remaining_defs
+
+        return schema
+
+    def _replace_refs_with_external(
+        self: Self,
+        schema: JsonSchema,
+        definitions: JsonSchemaDefinitions,
+        ref_template: str,
+    ) -> JsonSchema:
+        """Replace internal $ref with external references.
+
+        Only replaces refs for models that have enable_ref=True.
+
+        Args:
+            schema: The schema to process.
+            definitions: Dictionary of definitions to replace.
+            ref_template: Template for external references.
+
+        Returns:
+            Updated schema with external references.
+        """
+
+        def process_value(value: JsonValue) -> JsonValue:
+            if isinstance(value, dict):
+                if "$ref" in value:
+                    # Extract the definition name from the ref
+                    ref = value["$ref"]
+                    if ref.startswith(("#/$defs/", "#/definitions/")):
+                        def_name = ref.split("/")[-1]
+
+                        # Try to find the model info for this definition
+                        model_info = self._find_model_for_definition(def_name)
+
+                        if model_info:
+                            model_name, model_version = model_info
+
+                            # Check if this model is enabled for $ref
+                            if self.registry.is_ref_enabled(model_name, model_version):
+                                # Replace with external reference
+                                return {
+                                    "$ref": ref_template.format(
+                                        model=model_name, version=str(model_version)
+                                    )
+                                }
+                            # Keep as internal reference (will be inlined)
+                            return value
+
+                # Recursively process nested dictionaries
+                return {k: process_value(v) for k, v in value.items()}
+            if isinstance(value, list):
+                return [process_value(item) for item in value]
+            return value
+
+        return process_value(schema)  # type: ignore[return-value]
+
+    def _get_remaining_defs(
+        self: Self,
+        schema: JsonSchema,
+        original_defs: JsonSchemaDefinitions,
+    ) -> JsonSchemaDefinitions:
+        """Get definitions that should remain inline.
+
+        Args:
+            schema: The processed schema.
+            original_defs: Original definitions.
+
+        Returns:
+            Dictionary of definitions that weren't converted to external refs.
+        """
+        # Find all internal refs still in the schema
+        internal_refs: set[str] = set()
+
+        def find_internal_refs(value: dict[str, Any] | list[Any]) -> None:
+            if isinstance(value, dict):
+                if "$ref" in value:
+                    ref = value["$ref"]
+                    if ref.startswith(("#/$defs/", "#/definitions/")):
+                        def_name = ref.split("/")[-1]
+                        internal_refs.add(def_name)
+                for v in value.values():
+                    find_internal_refs(v)
+            elif isinstance(value, list):
+                for item in value:
+                    find_internal_refs(item)
+
+        find_internal_refs(schema)
+
+        # Return only definitions that are still referenced internally
+        return {k: v for k, v in original_defs.items() if k in internal_refs}
+
+    def _find_model_for_definition(self: Self, def_name: str) -> ModelMetadata | None:
+        """Find the registered model corresponding to a definition name.
+
+        Args:
+            def_name: The definition name from the schema.
+
+        Returns:
+            Tuple of (model_name, version) if found, None otherwise.
+        """
+        # Search through all registered models to find matching class name
+        for name, versions in self.registry._models.items():
+            for version, model_class in versions.items():
+                if model_class.__name__ == def_name:
+                    return (name, version)
+        return None
+
+    def get_all_schemas(self: Self, name: ModelName) -> dict[ModelVersion, JsonSchema]:
+        """Get all schemas for a model across all versions.
+
+        Args:
+            name: Name of the model.
+
+        Returns:
+            Dictionary mapping versions to their schemas.
+
+        Raises:
+            ValueError: If model not found.
+        """
+        if name not in self.registry._models:
+            raise ValueError(f"Model {name} not found")
+
+        return {
+            version: self.get_schema(name, version)
+            for version in self.registry._models[name]
+        }
+
+    def dump_schemas(
+        self: Self,
+        output_dir: str | Path,
+        indent: int = 2,
+        separate_definitions: bool = False,
+        ref_template: str | None = None,
+    ) -> None:
+        """Dump all schemas to JSON files.
+
+        Args:
+            output_dir: Directory path for output files.
+            indent: JSON indentation level.
+            separate_definitions: If True, create separate schema files for
+                nested models that have enable_ref=True.
+            ref_template: Template for $ref URLs when separate_definitions=True.
+                Defaults to relative file references if not provided.
+
+        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"
+            ... )
+        """
+        output_path = Path(output_dir)
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        if not separate_definitions:
+            # Original behavior: inline definitions
+            for name in self.registry._models:
+                for version, schema in self.get_all_schemas(name).items():
+                    file_path = output_path / f"{name}_v{version}.json"
+                    with open(file_path, "w", encoding="utf-8") as f:
+                        json.dump(schema, f, indent=indent)
+        else:
+            # New behavior: separate definition files for enable_ref=True models
+            # Default to relative file references
+            if ref_template is None:
+                ref_template = "{model}_v{version}.json"
+
+            # First pass: write all schemas
+            for name in self.registry._models:
+                for version in self.registry._models[name]:
+                    schema = self.get_schema_with_separate_defs(
+                        name, version, ref_template
+                    )
+                    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)
+
+    def get_nested_models(
+        self: Self,
+        name: ModelName,
+        version: str | ModelVersion,
+    ) -> list[ModelMetadata]:
+        """Get all nested models referenced by a model.
+
+        Args:
+            name: Name of the model.
+            version: Semantic version.
+
+        Returns:
+            List of (model_name, model_version) tuples for nested models.
+        """
+        ver = ModelVersion.parse(version) if isinstance(version, str) else version
+        model = self.registry.get_model(name, ver)
+
+        nested: list[ModelMetadata] = []
+
+        for field_info in model.model_fields.values():
+            model_type = self._get_model_type_from_field(field_info)
+            if model_type:
+                model_info = self.registry.get_model_info(model_type)
+                if model_info and model_info not in nested:
+                    nested.append(model_info)
+
+        return nested
+
+    def _get_model_type_from_field(
+        self: Self, field: FieldInfo
+    ) -> type[BaseModel] | None:
+        """Extract the Pydantic model type from a field.
+
+        Args:
+            field: The field info to extract from.
+
+        Returns:
+            The model type if found, None otherwise.
+        """
+        annotation = field.annotation
+        if annotation is None:
+            return None
+
+        # Handle direct model types
+        if isinstance(annotation, type) and issubclass(annotation, BaseModel):
+            return annotation
+
+        # Handle Optional, List, etc.
+        origin = get_origin(annotation)
+        if origin is not None:
+            args = get_args(annotation)
+            for arg in args:
+                if isinstance(arg, type) and issubclass(arg, BaseModel):
+                    return arg
+
+        return None

pyrmute/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

pyrmute/model_manager.py

@@ -0,0 +1,264 @@
+"""Model manager."""
+
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, Self
+
+from pydantic import BaseModel
+
+from ._migration_manager import MigrationManager
+from ._registry import Registry
+from ._schema_manager import SchemaManager
+from .model_version import ModelVersion
+from .types import (
+    DecoratedBaseModel,
+    JsonSchema,
+    JsonSchemaGenerator,
+    MigrationData,
+    MigrationFunc,
+    ModelMetadata,
+)
+
+
+class ModelManager:
+    """High-level interface for versioned model management.
+
+    Provides a unified API for model registration, migration, and schema
+    management.
+
+    Attributes:
+        registry: Registry instance.
+        migration_manager: MigrationManager instance.
+        schema_manager: SchemaManager instance.
+
+    Example:
+        >>> manager = ModelManager()
+        >>>
+        >>> @manager.model("User", "1.0.0")
+        ... class UserV1(BaseModel):
+        ...     name: str
+        >>>
+        >>> @manager.model("User", "2.0.0")
+        ... class UserV2(BaseModel):
+        ...     name: str
+        ...     email: str
+        >>>
+        >>> @manager.migration("User", "1.0.0", "2.0.0")
+        ... def migrate(data: MigrationData) -> MigrationData:
+        ...     return {**data, "email": "unknown@example.com"}
+    """
+
+    def __init__(self: Self) -> None:
+        """Initialize the versioned model manager."""
+        self.registry = Registry()
+        self.migration_manager = MigrationManager(self.registry)
+        self.schema_manager = SchemaManager(self.registry)
+
+    def model(
+        self: Self,
+        name: str,
+        version: str | ModelVersion,
+        schema_generator: JsonSchemaGenerator | None = None,
+        enable_ref: bool = False,
+    ) -> Callable[[type[DecoratedBaseModel]], type[DecoratedBaseModel]]:
+        """Register a versioned model.
+
+        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.
+
+        Returns:
+            Decorator function for model class.
+
+        Example:
+            >>> # Model that will be inlined (default)
+            >>> @manager.model("Address", "1.0.0")
+            ... class AddressV1(BaseModel):
+            ...     street: str
+            >>>
+            >>> # Model that can be a separate schema with $ref
+            >>> @manager.model("City", "1.0.0", enable_ref=True)
+            ... class CityV1(BaseModel):
+            ...     city: City
+        """
+        return self.registry.register(name, version, schema_generator, enable_ref)
+
+    def migration(
+        self: Self,
+        name: str,
+        from_version: str | ModelVersion,
+        to_version: str | ModelVersion,
+    ) -> Callable[[MigrationFunc], MigrationFunc]:
+        """Register a migration function.
+
+        Args:
+            name: Name of the model.
+            from_version: Source version.
+            to_version: Target version.
+
+        Returns:
+            Decorator function for migration function.
+        """
+        return self.migration_manager.register_migration(name, from_version, to_version)
+
+    def get(
+        self: Self, name: str, version: str | ModelVersion | None = None
+    ) -> type[BaseModel]:
+        """Get a model by name and version.
+
+        Args:
+            name: Name of the model.
+            version: Semantic version (returns latest if None).
+
+        Returns:
+            Model class.
+        """
+        if version is None:
+            return self.registry.get_latest(name)
+        return self.registry.get_model(name, version)
+
+    def migrate(
+        self: Self,
+        data: MigrationData,
+        name: str,
+        from_version: str | ModelVersion,
+        to_version: str | ModelVersion,
+    ) -> BaseModel:
+        """Migrate data between versions.
+
+        Args:
+            data: Data dictionary or BaseModel to migrate.
+            name: Name of the model.
+            from_version: Source version.
+            to_version: Target version.
+
+        Returns:
+            Migrated BaseModel.
+        """
+        migrated_data = self.migration_manager.migrate(
+            data, name, from_version, to_version
+        )
+        target_model = self.get(name, to_version)
+        return target_model.model_validate(migrated_data)
+
+    def get_schema(
+        self: Self,
+        name: str,
+        version: str | ModelVersion,
+        **kwargs: Any,
+    ) -> JsonSchema:
+        """Get JSON schema for a specific version.
+
+        Args:
+            name: Name of the model.
+            version: Semantic version.
+            **kwargs: Additional schema generation arguments.
+
+        Returns:
+            JSON schema dictionary.
+        """
+        return self.schema_manager.get_schema(name, version, **kwargs)
+
+    def list_models(self: Self) -> list[str]:
+        """Get list of all registered models.
+
+        Returns:
+            List of model names.
+        """
+        return self.registry.list_models()
+
+    def list_versions(self: Self, name: str) -> list[ModelVersion]:
+        """Get all versions for a model.
+
+        Args:
+            name: Name of the model.
+
+        Returns:
+            Sorted list of versions.
+        """
+        return self.registry.get_versions(name)
+
+    def dump_schemas(
+        self: Self,
+        output_dir: str | Path,
+        indent: int = 2,
+        separate_definitions: bool = False,
+        ref_template: str | None = None,
+    ) -> None:
+        """Export all schemas to JSON files.
+
+        Args:
+            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.
+            ref_template: Template for $ref URLs when separate_definitions=True.
+                Defaults to relative file references if not provided.
+
+        Example:
+            >>> # Inline definitions (default)
+            >>> manager.dump_schemas("schemas/")
+            >>>
+            >>> # Separate sub-schemas with relative refs
+            >>> 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"
+            ... )
+        """
+        self.schema_manager.dump_schemas(
+            output_dir, indent, separate_definitions, ref_template
+        )
+
+    def dump_schemas_with_refs(
+        self: Self,
+        output_dir: str | Path,
+        ref_template: str | None = None,
+        indent: int = 2,
+    ) -> None:
+        """Export schemas with separate files for nested models.
+
+        This is a convenience method that calls dump_schemas with
+        separate_definitions=True.
+
+        Args:
+            output_dir: Directory path for output.
+            ref_template: Template for $ref URLs. Supports {model} and
+                {version} placeholders. Defaults to relative file refs.
+            indent: JSON indentation level.
+
+        Example:
+            >>> # Relative file references (default)
+            >>> manager.dump_schemas_with_refs("schemas/")
+            >>>
+            >>> # Absolute URL references
+            >>> manager.dump_schemas_with_refs(
+            ...     "schemas/",
+            ...     ref_template="https://example.com/schemas/{model}_v{version}.json"
+            ... )
+        """
+        self.schema_manager.dump_schemas(
+            output_dir, indent, separate_definitions=True, ref_template=ref_template
+        )
+
+    def get_nested_models(
+        self: Self,
+        name: str,
+        version: str | ModelVersion,
+    ) -> list[ModelMetadata]:
+        """Get all nested models used by a model.
+
+        Args:
+            name: Name of the model.
+            version: Semantic version.
+
+        Returns:
+            List of (model_name, version) tuples for nested models.
+        """
+        return self.schema_manager.get_nested_models(name, version)

pyrmute/model_version.py

@@ -0,0 +1,60 @@
+"""Models the version provided to managers."""
+
+from dataclasses import dataclass
+from typing import Self
+
+
+@dataclass(frozen=True, order=True)
+class ModelVersion:
+    """Semantic version representation.
+
+    Attributes:
+        major: Major version number (breaking changes).
+        minor: Minor version number (backward-compatible features).
+        patch: Patch version number (backward-compatible fixes).
+    """
+
+    major: int
+    minor: int
+    patch: int
+
+    @classmethod
+    def parse(cls, version_str: str) -> Self:
+        """Parse a semantic version string.
+
+        Args:
+            version_str: Version string in format "major.minor.patch".
+
+        Returns:
+            Parsed Version instance.
+
+        Raises:
+            ValueError: If version string format is invalid.
+        """
+        parts = version_str.split(".")
+        if len(parts) != 3:  # noqa: PLR2004
+            raise ValueError(f"Invalid version format: {version_str}")
+
+        try:
+            major, minor, patch = int(parts[0]), int(parts[1]), int(parts[2])
+            if major < 0 or minor < 0 or patch < 0:
+                raise ValueError(f"Invalid version format: {version_str}")
+            return cls(major, minor, patch)
+        except ValueError as e:
+            raise ValueError(f"Invalid version format: {version_str}") from e
+
+    def __str__(self: Self) -> str:
+        """Return string representation of version.
+
+        Returns:
+            Version string in format "major.minor.patch".
+        """
+        return f"{self.major}.{self.minor}.{self.patch}"
+
+    def __repr__(self: Self) -> str:
+        """Return detailed string representation.
+
+        Returns:
+            Detailed version representation.
+        """
+        return f"ModelVersion({self.major}, {self.minor}, {self.patch})"