metaxy 0.0.1.dev3__py3-none-any.whl

metaxy/entrypoints.py

@@ -0,0 +1,296 @@
+"""Entrypoint discovery and loading for Metaxy features.
+
+This module provides functionality to automatically discover and load Feature
+classes from modules, supporting both:
+- Config-based entrypoints (list of module paths)
+- Environment-based entrypoints (environment variables starting with METAXY_ENTRYPOINT)
+- Package-based entrypoints (via importlib.metadata)
+
+Features are automatically registered to the active FeatureGraph when their
+containing modules are imported (via the Feature metaclass).
+"""
+
+import importlib
+import os
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from metaxy.models.feature import FeatureGraph
+
+from importlib.metadata import entry_points  # type: ignore[import-not-found]
+
+# Default entry point group name for package-based discovery
+DEFAULT_ENTRY_POINT_GROUP = "metaxy.project"
+
+
+class EntrypointLoadError(Exception):
+    """Raised when an entrypoint fails to load."""
+
+    pass
+
+
+def load_module_entrypoint(
+    module_path: str,
+    *,
+    graph: "FeatureGraph | None" = None,
+) -> None:
+    """Load a single module entrypoint.
+
+    Imports the specified module, which should contain Feature class definitions.
+    Features will be automatically registered to the active graph via the
+    Feature metaclass.
+
+    Args:
+        module_path: Fully qualified module path (e.g., "myapp.features.video")
+        graph: Target graph. If None, uses FeatureGraph.get_active()
+
+    Raises:
+        EntrypointLoadError: If module import fails
+
+    Example:
+        ```py
+        from metaxy.entrypoints import load_module_entrypoint
+        load_module_entrypoint("myapp.features.core")
+        # Features from myapp.features.core are now registered
+        ```
+    """
+    from metaxy.models.feature import FeatureGraph
+
+    target_graph = graph or FeatureGraph.get_active()
+
+    try:
+        # Set graph as active during import so Features register to it
+        with target_graph.use():
+            importlib.import_module(module_path)
+    except ImportError as e:
+        raise EntrypointLoadError(
+            f"Failed to import entrypoint module '{module_path}': {e}"
+        ) from e
+    except Exception as e:
+        raise EntrypointLoadError(
+            f"Error loading entrypoint module '{module_path}': {e}"
+        ) from e
+
+
+def load_entrypoints(
+    entrypoints: list[str],
+    *,
+    graph: "FeatureGraph | None" = None,
+) -> None:
+    """Load multiple module entrypoints from a list.
+
+    Args:
+        entrypoints: List of module paths to import
+        graph: Target graph. If None, uses FeatureGraph.get_active()
+
+    Raises:
+        EntrypointLoadError: If any module import fails
+
+    Example:
+        ```py
+        from metaxy.entrypoints import load_config_entrypoints
+        load_config_entrypoints([
+            "myapp.features.video",
+            "myapp.features.audio",
+            "myapp.features.text"
+        ])
+        ```
+    """
+    from metaxy.models.feature import FeatureGraph
+
+    target_graph = graph or FeatureGraph.get_active()
+
+    for module_path in entrypoints:
+        load_module_entrypoint(module_path, graph=target_graph)
+
+
+def load_package_entrypoints(
+    group: str = DEFAULT_ENTRY_POINT_GROUP,
+    *,
+    graph: "FeatureGraph | None" = None,
+) -> None:
+    """Load entrypoints from installed packages using importlib.metadata.
+
+    Discovers and loads all entry points registered in the specified group.
+    This is the package-based entrypoint mechanism using standard Python
+    packaging infrastructure.
+
+    Packages declare entrypoints in their pyproject.toml:
+        [project.entry-points."metaxy.project"]
+        my-project = "mypackage:init"
+        # or point directly to a module
+        my-project = "mypackage.features"
+
+    The entry point can reference either:
+    - A callable function (module:function syntax) that will be invoked to load features
+    - A module (module syntax) that contains Feature definitions (importing registers them)
+
+    Args:
+        group: Entry point group name (default: "metaxy.project")
+        graph: Target graph. If None, uses FeatureGraph.get_active()
+
+    Raises:
+        EntrypointLoadError: If any entrypoint fails to load
+
+    Example:
+        ```py
+        from metaxy.entrypoints import load_package_entrypoints
+        # Discover and load all installed plugins
+        load_package_entrypoints()
+        ```
+    """
+    from metaxy.models.feature import FeatureGraph
+
+    target_graph = graph or FeatureGraph.get_active()
+
+    # Discover entry points
+    # Note: Python 3.10+ returns SelectableGroups, 3.9 returns dict
+    discovered = entry_points()
+
+    # Handle different return types across Python versions
+    if hasattr(discovered, "select"):
+        # Python 3.10+: SelectableGroups with select() method
+        eps = discovered.select(group=group)
+    else:
+        # Python 3.9: dict-like interface
+        eps = discovered.get(group, [])
+
+    eps_list = list(eps)
+
+    if not eps_list:
+        return
+
+    for ep in eps_list:
+        try:
+            # Load the entry point (imports the module and returns the object)
+            with target_graph.use():
+                loaded = ep.load()
+                # If it's callable (module:function syntax), call it
+                # If it's a module (just module syntax), importing already registered features
+                if callable(loaded):
+                    loaded()
+        except Exception as e:
+            raise EntrypointLoadError(
+                f"Failed to load package entrypoint '{ep.name}' ({ep.value}): {e}"
+            ) from e
+
+
+def load_env_entrypoints() -> None:
+    """Load entrypoints from environment variables.
+
+    Discovers and loads all entry points from environment variables matching
+    the pattern METAXY_ENTRYPOINT*. Each variable should contain a
+    comma-separated list of module paths.
+
+    Environment variables:
+        METAXY_ENTRYPOINT="myapp.features.core,myapp.features.extra"
+        METAXY_ENTRYPOINT_PLUGINS="plugin1.features,plugin2.features"
+
+    Args:
+        graph: Target graph. If None, uses FeatureGraph.get_active()
+
+    Raises:
+        EntrypointLoadError: If any entrypoint fails to load
+
+    Example:
+        ```py
+        import os
+        os.environ["METAXY_ENTRYPOINT"] = "myapp.features.core"
+        from metaxy.entrypoints import load_env_entrypoints
+        load_env_entrypoints()
+        ```
+    """
+
+    # Find all environment variables matching METAXY_ENTRYPOINT*
+    env_vars = {
+        key: value
+        for key, value in os.environ.items()
+        if key.startswith("METAXY_ENTRYPOINT")
+    }
+
+    if not env_vars:
+        return
+
+    # Collect all module paths from all matching env vars
+    all_module_paths = []
+    for env_var, value in sorted(env_vars.items()):
+        # Split by comma and strip whitespace
+        module_paths = [path.strip() for path in value.split(",") if path.strip()]
+        all_module_paths.extend(module_paths)
+
+    if not all_module_paths:
+        return
+
+    # Load each module path
+    for module_path in all_module_paths:
+        load_module_entrypoint(module_path)
+
+
+def load_features(
+    entrypoints: list[str] | None = None,
+    package_entrypoint_group: str = DEFAULT_ENTRY_POINT_GROUP,
+    *,
+    load_config: bool = True,
+    load_packages: bool = True,
+    load_env: bool = True,
+) -> "FeatureGraph":
+    """Discover and load all entrypoints from config, packages, and environment.
+
+    This is the main entry point for loading features. It combines config-based,
+    package-based, and environment-based entrypoint discovery.
+
+    Args:
+        entrypoints: List of module paths (optional)
+        package_entrypoint_group: Entry point group for package discovery
+        load_config: Whether to load config-based entrypoints (default: True)
+        load_packages: Whether to load package-based entrypoints (default: True)
+        load_env: Whether to load environment-based entrypoints (default: True)
+
+    Returns:
+        The graph that was populated (useful for testing/inspection)
+
+    Raises:
+        EntrypointLoadError: If any entrypoint fails to load
+
+    Example:
+        ```py
+        from metaxy.entrypoints import load_features
+
+        # Load from all sources
+        graph = load_features(
+            entrypoints=["myapp.features.core"],
+            load_packages=True,
+            load_env=True
+        )
+
+        # Load only from config
+        graph = load_features(
+            entrypoints=["myapp.features.core"],
+            load_packages=False,
+            load_env=False
+        )
+        ```
+    """
+    from metaxy.config import MetaxyConfig
+    from metaxy.models.feature import FeatureGraph
+
+    target_graph = FeatureGraph.get_active()
+
+    # Load explicit entrypoints
+    if entrypoints:
+        load_entrypoints(entrypoints)
+
+    # Load config-based entrypoints
+    if load_config:
+        config = MetaxyConfig.load(search_parents=True)
+        load_entrypoints(config.entrypoints)
+
+    # Load package-based entrypoints
+    if load_packages:
+        load_package_entrypoints(package_entrypoint_group)
+
+    # Load environment-based entrypoints
+    if load_env:
+        load_env_entrypoints()
+
+    return target_graph

metaxy/ext/__init__.py

@@ -0,0 +1 @@
+"""Integrations with third-party software."""

metaxy/ext/dagster/__init__.py

@@ -0,0 +1,54 @@
+from metaxy.ext.dagster.constants import (
+    DAGSTER_METAXY_FEATURE_METADATA_KEY,
+    DAGSTER_METAXY_INFO_METADATA_KEY,
+    DAGSTER_METAXY_KIND,
+    DAGSTER_METAXY_PARTITION_KEY,
+    DAGSTER_METAXY_PROJECT_TAG_KEY,
+    METAXY_DAGSTER_METADATA_KEY,
+)
+from metaxy.ext.dagster.dagster_type import feature_to_dagster_type
+from metaxy.ext.dagster.io_manager import MetaxyIOManager, MetaxyOutput
+from metaxy.ext.dagster.metaxify import metaxify
+from metaxy.ext.dagster.observable import observable_metaxy_asset
+from metaxy.ext.dagster.resources import MetaxyStoreFromConfigResource
+from metaxy.ext.dagster.selection import select_metaxy_assets
+from metaxy.ext.dagster.table_metadata import (
+    build_column_lineage,
+    build_column_schema,
+    build_table_preview_metadata,
+)
+from metaxy.ext.dagster.utils import (
+    FeatureStats,
+    build_partition_filter,
+    compute_feature_stats,
+    compute_stats_from_lazy_frame,
+    generate_materialize_results,
+    generate_observe_results,
+    get_partition_filter,
+)
+
+__all__ = [
+    "metaxify",
+    "feature_to_dagster_type",
+    "build_column_schema",
+    "build_column_lineage",
+    "build_table_preview_metadata",
+    "observable_metaxy_asset",
+    "select_metaxy_assets",
+    "generate_materialize_results",
+    "generate_observe_results",
+    "compute_feature_stats",
+    "compute_stats_from_lazy_frame",
+    "get_partition_filter",
+    "build_partition_filter",
+    "FeatureStats",
+    "MetaxyStoreFromConfigResource",
+    "MetaxyIOManager",
+    "MetaxyOutput",
+    "METAXY_DAGSTER_METADATA_KEY",
+    "DAGSTER_METAXY_FEATURE_METADATA_KEY",
+    "DAGSTER_METAXY_INFO_METADATA_KEY",
+    "DAGSTER_METAXY_KIND",
+    "DAGSTER_METAXY_PARTITION_KEY",
+    "DAGSTER_METAXY_PROJECT_TAG_KEY",
+]

metaxy/ext/dagster/constants.py

@@ -0,0 +1,10 @@
+DAGSTER_METAXY_FEATURE_METADATA_KEY = "metaxy/feature"
+DAGSTER_METAXY_KIND = "metaxy"
+DAGSTER_METAXY_PARTITION_KEY = "partition_by"
+DAGSTER_METAXY_PROJECT_TAG_KEY = "metaxy/project"
+
+DAGSTER_METAXY_INFO_METADATA_KEY = "metaxy/info"
+DAGSTER_COLUMN_SCHEMA_METADATA_KEY = "dagster/column_schema"
+DAGSTER_COLUMN_LINEAGE_METADATA_KEY = "dagster/column_lineage"
+
+METAXY_DAGSTER_METADATA_KEY = "dagster/attributes"

metaxy/ext/dagster/dagster_type.py

@@ -0,0 +1,156 @@
+"""DagsterType builder for Metaxy features.
+
+This module provides utilities for creating Dagster types that validate
+Metaxy feature outputs with proper metadata injection (table schema, etc.).
+"""
+
+from collections.abc import Callable, Mapping
+from typing import Any
+
+import dagster as dg
+import narwhals as nw
+
+import metaxy as mx
+from metaxy.ext.dagster.constants import (
+    DAGSTER_COLUMN_LINEAGE_METADATA_KEY,
+    DAGSTER_COLUMN_SCHEMA_METADATA_KEY,
+    DAGSTER_METAXY_INFO_METADATA_KEY,
+)
+from metaxy.ext.dagster.table_metadata import build_column_lineage, build_column_schema
+from metaxy.ext.dagster.utils import build_feature_info_metadata
+
+
+def _create_type_check_fn(
+    feature_key: mx.FeatureKey,
+) -> Callable[[dg.TypeCheckContext, Any], dg.TypeCheck]:
+    """Create a type check function for a Metaxy feature.
+
+    The type check function validates that the output is either:
+    - None (allowed for MetaxyOutput)
+    - A narwhals-compatible dataframe (IntoFrame)
+
+    Args:
+        feature_key: The Metaxy feature key for error messages.
+
+    Returns:
+        A callable type check function for DagsterType.
+    """
+
+    def type_check_fn(context: dg.TypeCheckContext, value: Any) -> dg.TypeCheck:
+        # None is a valid MetaxyOutput (indicates no data to write)
+        if value is None:
+            return dg.TypeCheck(success=True)
+
+        # Try to convert to narwhals frame - this validates the type
+        try:
+            nw.from_native(value)
+            return dg.TypeCheck(success=True)
+        except TypeError as e:
+            return dg.TypeCheck(
+                success=False,
+                description=(
+                    f"Expected a narwhals-compatible dataframe or None for "
+                    f"Metaxy feature '{feature_key.to_string()}', "
+                    f"but got {type(value).__name__}:\n{e}"
+                ),
+            )
+
+    return type_check_fn
+
+
+def feature_to_dagster_type(
+    feature: mx.CoercibleToFeatureKey,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    inject_column_schema: bool = True,
+    inject_column_lineage: bool = True,
+    metadata: Mapping[str, Any] | None = None,
+) -> dg.DagsterType:
+    """Build a Dagster type from a Metaxy feature.
+
+    Creates a `dagster.DagsterType` that validates outputs are
+    [`MetaxyOutput`][metaxy.ext.dagster.MetaxyOutput] (i.e., narwhals-compatible
+    dataframes or `None`) and includes metadata derived from the feature's Pydantic
+    model fields.
+
+    Args:
+        feature: The Metaxy feature to create a type for. Can be a feature class,
+            feature key, or string that can be coerced to a feature key.
+        name: Optional custom name for the DagsterType. Defaults to the feature's
+            table name (e.g., "project__feature_name").
+        description: Optional custom description. Defaults to the feature class
+            docstring or a generated description.
+        inject_column_schema: Whether to inject the column schema as metadata.
+            The schema is derived from Pydantic model fields.
+        inject_column_lineage: Whether to inject column lineage as metadata.
+            The lineage is derived from feature dependencies.
+        metadata: Optional custom metadata to inject into the DagsterType.
+
+    Returns:
+        A DagsterType configured for the Metaxy feature with appropriate
+        type checking and metadata.
+
+    !!! tip
+        This is automatically injected by [`@metaxify`][metaxy.ext.dagster.metaxify.metaxify]
+
+    Example:
+        ```python
+        import dagster as dg
+        import polars as pl
+        import metaxy.ext.dagster as mxd
+        from myproject.features import MyFeature  # Your Metaxy feature class
+
+        @mxd.metaxify(feature=MyFeature)
+        @dg.asset(dagster_type=mxd.feature_to_dagster_type(MyFeature))
+        def my_asset():
+            return pl.DataFrame({"id": [1, 2, 3], "value": ["a", "b", "c"]})
+        ```
+
+    !!! info "See also"
+        - [`metaxify`][metaxy.ext.dagster.metaxify.metaxify]: Decorator for injecting
+          Metaxy metadata into Dagster assets.
+        - [`MetaxyOutput`][metaxy.ext.dagster.MetaxyOutput]: The type alias for valid
+          Metaxy outputs.
+    """
+    from metaxy.ext.dagster.io_manager import MetaxyOutput
+
+    feature_key = mx.coerce_to_feature_key(feature)
+    feature_cls = mx.get_feature_by_key(feature_key)
+
+    # Determine name
+    type_name = name or feature_key.table_name
+
+    # Determine description
+    if description is None:
+        if feature_cls.__doc__:
+            import inspect
+
+            description = inspect.cleandoc(feature_cls.__doc__)
+        else:
+            description = f"Metaxy feature '{feature_key.to_string()}'."
+
+    # Build metadata - start with custom metadata if provided
+    final_metadata: dict[str, Any] = dict(metadata) if metadata else {}
+    final_metadata[DAGSTER_METAXY_INFO_METADATA_KEY] = build_feature_info_metadata(
+        feature_key
+    )
+    if inject_column_schema:
+        column_schema = build_column_schema(feature_cls)
+        if column_schema is not None:
+            final_metadata[DAGSTER_COLUMN_SCHEMA_METADATA_KEY] = column_schema
+
+    if inject_column_lineage:
+        column_lineage = build_column_lineage(feature_cls)
+        if column_lineage is not None:
+            final_metadata[DAGSTER_COLUMN_LINEAGE_METADATA_KEY] = column_lineage
+
+    dagster_type = dg.DagsterType(
+        type_check_fn=_create_type_check_fn(feature_key),
+        name=type_name,
+        description=description,
+        typing_type=MetaxyOutput,
+        metadata=final_metadata,
+    )
+
+    return dagster_type