metaxy 0.0.1.dev3__py3-none-any.whl

metaxy/models/types.py

@@ -0,0 +1,475 @@
+"""Type definitions for metaxy models."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, TypeAlias, overload
+
+from pydantic import (
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    RootModel,
+    TypeAdapter,
+    field_validator,
+    model_serializer,
+    model_validator,
+)
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+KEY_SEPARATOR = "/"
+
+# backcompat
+FEATURE_KEY_SEPARATOR = KEY_SEPARATOR
+FIELD_KEY_SEPARATOR = KEY_SEPARATOR
+
+
+class SnapshotPushResult(NamedTuple):
+    """Result of recording a feature graph snapshot.
+
+    Attributes:
+        snapshot_version: The deterministic hash of the graph snapshot
+        already_pushed: True if this snapshot_version was already pushed previously
+        updated_features: List of feature keys with updated information (changed full_definition_version)
+    """
+
+    snapshot_version: str
+    already_pushed: bool
+    updated_features: list[str]
+
+
+_CoercibleToKey: TypeAlias = Sequence[str] | str
+
+
+class _Key(RootModel[tuple[str, ...]]):
+    """
+    A common class for key-like objects that contain a sequence of string parts.
+
+    Parts cannot contain forward slashes (/) or double underscores (__).
+
+    Args:
+        key: Feature key as string ("a/b/c"), sequence (["a", "b", "c"]), or FeatureKey instance.
+             String format is split on "/" separator.
+    """
+
+    model_config = ConfigDict(frozen=True, repr=False)  # pyright: ignore[reportCallIssue]  # Make immutable for hashability, use custom __repr__
+
+    root: tuple[str, ...]
+
+    if TYPE_CHECKING:
+
+        @overload
+        def __init__(self, parts: str) -> None: ...
+
+        @overload
+        def __init__(self, parts: Sequence[str]) -> None: ...
+
+        @overload
+        def __init__(self, parts: Self) -> None: ...
+
+        def __init__(  # pyright: ignore[reportMissingSuperCall]
+            self,
+            parts: str | Sequence[str] | Self,
+        ) -> None: ...
+
+    @model_validator(mode="before")
+    @classmethod
+    def _validate_input(cls, data: Any) -> Any:
+        """Convert various input types to tuple of strings."""
+        # If it's already a tuple, validate and return it
+        if isinstance(data, tuple):
+            return data
+
+        # Handle dict input (from Pydantic deserialization)
+        if isinstance(data, dict):
+            # RootModel deserialization passes dict with "root" key
+            if "root" in data:
+                root_value = data["root"]
+                if isinstance(root_value, tuple):
+                    return root_value
+                elif isinstance(root_value, (list, Sequence)):
+                    return tuple(root_value)
+            # Legacy "parts" key for backward compatibility
+            elif "parts" in data:
+                parts = data["parts"]
+                if (
+                    isinstance(parts, (list, tuple))
+                    and parts
+                    and isinstance(parts[0], dict)
+                ):
+                    # Handle incorrectly nested structure like {'parts': [{'parts': [...]}]}
+                    if "parts" in parts[0]:
+                        return tuple(parts[0]["parts"])
+                    else:
+                        raise ValueError(f"Invalid nested structure in parts: {parts}")
+                return tuple(parts) if not isinstance(parts, tuple) else parts
+            # Empty dict
+            raise ValueError("Dict must contain 'root' or 'parts' key")
+
+        # Handle string input - split on separator
+        if isinstance(data, str):
+            return tuple(data.split(KEY_SEPARATOR))
+
+        # Handle instance of same class - extract root
+        if isinstance(data, cls):
+            return data.root
+
+        # Handle sequence (list, etc.)
+        if isinstance(data, Sequence):
+            return tuple(data)
+
+        raise ValueError(f"Cannot create {cls.__name__} from {type(data).__name__}")
+
+    @field_validator("root", mode="after")
+    @classmethod
+    def _validate_root_content(cls, value: tuple[str, ...]) -> tuple[str, ...]:
+        """Validate that parts don't contain forbidden characters."""
+        for part in value:
+            if not isinstance(part, str):
+                raise ValueError(
+                    f"{cls.__name__} parts must be strings, got {type(part).__name__}"
+                )
+            if "/" in part:
+                raise ValueError(
+                    f"{cls.__name__} part '{part}' cannot contain forward slashes (/). "
+                    f"Forward slashes are reserved as the separator in to_string(). "
+                    f"Use underscores or hyphens instead."
+                )
+            if "__" in part:
+                raise ValueError(
+                    f"{cls.__name__} part '{part}' cannot contain double underscores (__). "
+                    f"Use single underscores or hyphens instead."
+                )
+        return value
+
+    @model_serializer
+    def _serialize_model(self) -> list[str]:
+        """Serialize to list format for backward compatibility."""
+        return list(self.parts)
+
+    @property
+    def parts(self) -> tuple[str, ...]:
+        """Backward compatibility property for accessing root as parts."""
+        return self.root
+
+    def to_string(self) -> str:
+        """Convert to string representation with "/" separator."""
+        return KEY_SEPARATOR.join(self.parts)
+
+    def to_struct_key(self) -> str:
+        """Convert to a name that can be used as struct key in databases"""
+        return "_".join(self.parts)
+
+    def to_column_suffix(self) -> str:
+        """Convert to a suffix usable for database column names (typically temporary)."""
+        return "__" + "_".join(self.parts)
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return self.to_string()
+
+    def __str__(self) -> str:
+        """Return string representation."""
+        return self.to_string()
+
+    def __lt__(self, other: Any) -> bool:
+        """Less than comparison for sorting."""
+        if isinstance(other, self.__class__):
+            return self.parts < other.parts
+        return NotImplemented
+
+    def __le__(self, other: Any) -> bool:
+        """Less than or equal comparison for sorting."""
+        if isinstance(other, self.__class__):
+            return self.parts <= other.parts
+        return NotImplemented
+
+    def __gt__(self, other: Any) -> bool:
+        """Greater than comparison for sorting."""
+        if isinstance(other, self.__class__):
+            return self.parts > other.parts
+        return NotImplemented
+
+    def __ge__(self, other: Any) -> bool:
+        """Greater than or equal comparison for sorting."""
+        if isinstance(other, self.__class__):
+            return self.parts >= other.parts
+        return NotImplemented
+
+    def __iter__(self) -> Iterator[str]:  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Return iterator over parts."""
+        return iter(self.parts)
+
+    @property
+    def table_name(self) -> str:
+        """Get SQL-like table name for this feature key.
+
+        Replaces hyphens with underscores for SQL compatibility.
+        """
+        return "__".join(part.replace("-", "_") for part in self.parts)
+
+    # List-like interface for backward compatibility
+    def __getitem__(self, index: int) -> str:
+        """Get part by index."""
+        return self.parts[index]
+
+    def __len__(self) -> int:
+        """Get number of parts."""
+        return len(self.parts)
+
+    def __contains__(self, item: str) -> bool:
+        """Check if part is in key."""
+        return item in self.parts
+
+    def __reversed__(self):
+        """Return reversed iterator over parts."""
+        return reversed(self.parts)
+
+
+# CoercibleToKey: TypeAlias = _CoercibleToKey | _Key
+
+
+class FeatureKey(_Key):
+    """
+    Feature key as a sequence of string parts.
+
+    Hashable for use as dict keys in registries.
+    Parts cannot contain forward slashes (/) or double underscores (__).
+
+    Example:
+
+        ```py
+        FeatureKey("a/b/c")  # String format
+        # FeatureKey(parts=['a', 'b', 'c'])
+
+        FeatureKey(["a", "b", "c"])  # List format
+        # FeatureKey(parts=['a', 'b', 'c'])
+
+        FeatureKey(FeatureKey(["a", "b", "c"]))  # FeatureKey copy
+        # FeatureKey(parts=['a', 'b', 'c'])
+        ```
+    """
+
+    if TYPE_CHECKING:
+
+        @overload
+        def __init__(self, parts: str) -> None: ...
+
+        @overload
+        def __init__(self, parts: Sequence[str]) -> None: ...
+
+        @overload
+        def __init__(self, parts: FeatureKey) -> None: ...
+
+        def __init__(  # pyright: ignore[reportMissingSuperCall]
+            self,
+            parts: str | Sequence[str] | FeatureKey,
+        ) -> None: ...
+
+    def model_dump(self, **kwargs: Any) -> Any:
+        """Serialize to list format for backward compatibility."""
+        # When serializing this key, return it as a list of parts
+        # instead of the full Pydantic model structure
+        return list(self.parts)
+
+    def __hash__(self) -> int:
+        """Return hash for use as dict keys."""
+        return hash(self.parts)
+
+    def __eq__(self, other: Any) -> bool:
+        """Check equality with another instance."""
+        if isinstance(other, self.__class__):
+            return self.parts == other.parts
+        return super().__eq__(other)
+
+    def to_column_suffix(self) -> str:
+        """Convert to a suffix usable for database column names (typically temporary)."""
+        return "__" + "_".join(self.parts)
+
+
+class FieldKey(_Key):
+    """
+    Field key as a sequence of string parts.
+
+    Hashable for use as dict keys in registries.
+    Parts cannot contain forward slashes (/) or double underscores (__).
+
+    Example:
+
+        ```py
+        FieldKey("a/b/c")  # String format
+        # FieldKey(parts=['a', 'b', 'c'])
+
+        FieldKey(["a", "b", "c"])  # List format
+        # FieldKey(parts=['a', 'b', 'c'])
+
+        FieldKey(FieldKey(["a", "b", "c"]))  # FieldKey copy
+        # FieldKey(parts=['a', 'b', 'c'])
+        ```
+    """
+
+    if TYPE_CHECKING:
+
+        @overload
+        def __init__(self, parts: str) -> None: ...
+
+        @overload
+        def __init__(self, parts: Sequence[str]) -> None: ...
+
+        @overload
+        def __init__(self, parts: FieldKey) -> None: ...
+
+        def __init__(  # pyright: ignore[reportMissingSuperCall]
+            self,
+            parts: str | Sequence[str] | FieldKey,
+        ) -> None: ...
+
+    def model_dump(self, **kwargs: Any) -> Any:
+        """Serialize to list format for backward compatibility."""
+        # When serializing this key, return it as a list of parts
+        # instead of the full Pydantic model structure
+        return list(self.parts)
+
+    def __hash__(self) -> int:
+        """Return hash for use as dict keys."""
+        return hash(self.parts)
+
+    def __eq__(self, other: Any) -> bool:
+        """Check equality with another instance."""
+        if isinstance(other, self.__class__):
+            return self.parts == other.parts
+        return super().__eq__(other)
+
+
+_CoercibleToFeatureKey: TypeAlias = _CoercibleToKey | FeatureKey
+
+FeatureKeyAdapter = TypeAdapter(
+    FeatureKey
+)  # can call .validate_python() to transform acceptable types into a FeatureKey
+FieldKeyAdapter = TypeAdapter(
+    FieldKey
+)  # can call .validate_python() to transform acceptable types into a FieldKey
+
+
+def _coerce_to_feature_key(value: Any) -> FeatureKey:
+    """Convert various types to FeatureKey.
+
+    Accepts:
+
+    - slashed `str`: `"a/b/c"`
+
+    - `Sequence[str]`: `["a", "b", "c"]`
+
+    - `FeatureKey`: pass through
+
+    - `type[BaseFeature]`: extracts .spec().key
+
+    Args:
+        value: Value to coerce to `FeatureKey`
+
+    Returns:
+        canonical `FeatureKey` instance
+
+    Raises:
+        ValidationError: If value cannot be coerced to FeatureKey
+    """
+    if isinstance(value, FeatureKey):
+        return value
+
+    # Check if it's a BaseFeature class
+    # Import here to avoid circular dependency at module level
+    from metaxy.models.feature import BaseFeature
+
+    if isinstance(value, type) and issubclass(value, BaseFeature):
+        return value.spec().key
+
+    # Handle str, Sequence[str]
+    return FeatureKeyAdapter.validate_python(value)
+
+
+def _coerce_to_field_key(value: Any) -> FieldKey:
+    """Convert various types to FieldKey.
+
+    Accepts:
+
+        - slashed `str`: `"a/b/c"`
+
+        - `Sequence[str]`: `["a", "b", "c"]`
+
+        - `FieldKey`: pass through
+
+    Args:
+        value: Value to coerce to `FieldKey`
+
+    Returns:
+        canonical `FieldKey` instance
+
+    Raises:
+        ValidationError: If value cannot be coerced to `FieldKey`
+    """
+    if isinstance(value, FieldKey):
+        return value
+
+    # Handle str, Sequence[str]
+    return FieldKeyAdapter.validate_python(value)
+
+
+if TYPE_CHECKING:
+    from metaxy.models.feature import BaseFeature
+
+# Type unions - what inputs are accepted
+CoercibleToFeatureKey: TypeAlias = (
+    str | Sequence[str] | FeatureKey | type["BaseFeature"]
+)
+CoercibleToFieldKey: TypeAlias = str | Sequence[str] | FieldKey
+
+# Annotated types for Pydantic field annotations - automatically validate
+# After validation, these ARE FeatureKey/FieldKey (not unions)
+ValidatedFeatureKey: TypeAlias = Annotated[
+    FeatureKey,
+    BeforeValidator(_coerce_to_feature_key),
+    Field(
+        description="Feature key. Accepts a slashed string ('a/b/c'), a sequence of strings, a FeatureKey instance, or a child class of BaseFeature"
+    ),
+]
+
+ValidatedFieldKey: TypeAlias = Annotated[
+    FieldKey,
+    BeforeValidator(_coerce_to_field_key),
+    Field(
+        description="Field key. Accepts a slashed string ('a/b/c'), a sequence of strings, or a FieldKey instance."
+    ),
+]
+
+# TypeAdapters for non-Pydantic usage (e.g., in metadata_store/base.py)
+ValidatedFeatureKeyAdapter: TypeAdapter[ValidatedFeatureKey] = TypeAdapter(
+    ValidatedFeatureKey
+)
+ValidatedFieldKeyAdapter: TypeAdapter[ValidatedFieldKey] = TypeAdapter(
+    ValidatedFieldKey
+)
+
+
+# Collection types for common patterns - automatically validate sequences
+# Pydantic will validate each element using ValidatedFeatureKey/ValidatedFieldKey
+ValidatedFeatureKeySequence: TypeAlias = Annotated[
+    Sequence[ValidatedFeatureKey],
+    Field(description="Sequence items coerced into FeatureKey."),
+]
+
+ValidatedFieldKeySequence: TypeAlias = Annotated[
+    Sequence[ValidatedFieldKey],
+    Field(description="Sequence items coerced into FieldKey."),
+]
+
+# TypeAdapters for non-Pydantic usage
+ValidatedFeatureKeySequenceAdapter: TypeAdapter[ValidatedFeatureKeySequence] = (
+    TypeAdapter(ValidatedFeatureKeySequence)
+)
+ValidatedFieldKeySequenceAdapter: TypeAdapter[ValidatedFieldKeySequence] = TypeAdapter(
+    ValidatedFieldKeySequence
+)
+
+FeatureDepMetadata: TypeAlias = dict[str, Any]

metaxy/utils/__init__.py

@@ -0,0 +1 @@
+"""Utility modules for Metaxy."""

metaxy/utils/constants.py

@@ -0,0 +1,2 @@
+# Sentinel value for code_version in struct serialization
+DEFAULT_CODE_VERSION = "metaxy__intitial"

metaxy/utils/exceptions.py

@@ -0,0 +1,23 @@
+class MetaxyError(Exception):
+    """Base class for all errors thrown by the Metaxy framework.
+
+    Users should not subclass this base class for their own exceptions.
+    """
+
+    @property
+    def is_user_code_error(self):
+        """Returns true if this error is attributable to user code."""
+        return False
+
+
+class MetaxyInvariantViolationError(MetaxyError):
+    """Indicates the user has violated a well-defined invariant that can only be enforced
+    at runtime.
+    """
+
+
+class MetaxyEmptyCodeVersionError(MetaxyInvariantViolationError):
+    """Indicates that an empty code version was provided where it is not allowed.
+
+    Code version must be a non-empty string.
+    """