metaxy 0.0.0__py3-none-any.whl

metaxy/graph/diff/rendering/rich.py

@@ -0,0 +1,165 @@
+"""Terminal renderer using Rich Tree for hierarchical display.
+
+Requires rich library to be installed.
+"""
+
+from metaxy.graph.diff.models import NodeStatus
+from metaxy.graph.diff.rendering.base import BaseRenderer
+
+
+class TerminalRenderer(BaseRenderer):
+    """Renders graph using Rich Tree for terminal display.
+
+    Creates a hierarchical tree view with colors and icons.
+    Supports both normal and diff rendering via node status.
+    """
+
+    def render(self) -> str:
+        """Render graph as Rich Tree for terminal.
+
+        Returns:
+            Rendered tree as string with ANSI color codes
+        """
+        from rich.console import Console
+        from rich.tree import Tree
+
+        console = Console()
+
+        # Get filtered graph data based on config
+        filtered_graph = self._get_filtered_graph_data()
+
+        # Create root node
+        if self.config.show_snapshot_version:
+            snapshot_version = self._format_hash(filtered_graph.snapshot_version)
+            root = Tree(
+                f"📊 [bold]Graph[/bold] [dim](snapshot: {snapshot_version})[/dim]"
+            )
+        else:
+            root = Tree("📊 [bold]Graph[/bold]")
+
+        # Create walker for filtered graph and add features in topological order
+        from metaxy.graph.diff.traversal import GraphWalker
+
+        walker = GraphWalker(filtered_graph)
+        for node in walker.topological_sort():
+            self._render_feature_node(root, node)
+
+        # Render to string
+        with console.capture() as capture:
+            console.print(root)
+        return capture.get()
+
+    def _render_feature_node(self, parent, node):
+        """Add a feature node to the tree.
+
+        Args:
+            parent: Parent tree node
+            node: GraphNode
+        """
+        # Get status color
+        status_color = self._get_status_color(node.status)
+
+        # Build feature label
+        label_parts = [
+            f"[{status_color}]{self._format_feature_key(node.key)}[/{status_color}]"
+        ]
+
+        # Show version info
+        if self.config.show_feature_versions:
+            if node.status == NodeStatus.CHANGED and node.old_version is not None:
+                # Show version transition for changed nodes
+                version_transition = self._format_version_transition(
+                    node.old_version, node.version
+                )
+                label_parts.append(version_transition)
+            else:
+                # Normal version display
+                version = self._format_hash(node.version)
+                label_parts.append(f"[yellow](v: {version})[/yellow]")
+
+        if self.config.show_code_versions and node.code_version is not None:
+            label_parts.append(f"[dim](cv: {node.code_version})[/dim]")
+
+        # Add status badge for diff mode
+        if node.status != NodeStatus.NORMAL:
+            status_badge = self._get_status_badge(node.status)
+            label_parts.append(status_badge)
+
+        label = " ".join(label_parts)
+        feature_branch = parent.add(label)
+
+        # Add fields
+        if self.config.show_fields and node.fields:
+            fields_branch = feature_branch.add("🔧 [green]fields[/green]")
+            for field_node in node.fields:
+                self._render_field_node(fields_branch, field_node)
+
+        # Add dependencies
+        if node.dependencies:
+            deps_branch = feature_branch.add("⬅️  [blue]depends on[/blue]")
+            for dep_key in node.dependencies:
+                dep_color = status_color  # Use same color as parent for simplicity
+                deps_branch.add(
+                    f"[{dep_color}]{self._format_feature_key(dep_key)}[/{dep_color}]"
+                )
+
+    def _render_field_node(self, parent, field_node):
+        """Add a field node to the tree.
+
+        Args:
+            parent: Parent tree node
+            field_node: FieldNode
+        """
+        # Get status color
+        status_color = self._get_status_color(field_node.status)
+
+        label_parts = [
+            f"[{status_color}]{self._format_field_key(field_node.key)}[/{status_color}]"
+        ]
+
+        # Show version info
+        if self.config.show_field_versions:
+            if (
+                field_node.status == NodeStatus.CHANGED
+                and field_node.old_version is not None
+            ):
+                # Show version transition for changed fields
+                version_transition = self._format_version_transition(
+                    field_node.old_version, field_node.version
+                )
+                label_parts.append(version_transition)
+            else:
+                # Normal version display
+                version = self._format_hash(field_node.version)
+                label_parts.append(f"[yellow](v: {version})[/yellow]")
+
+        if self.config.show_code_versions and field_node.code_version is not None:
+            label_parts.append(f"[dim](cv: {field_node.code_version})[/dim]")
+
+        # Add status badge for diff mode
+        if field_node.status != NodeStatus.NORMAL:
+            status_badge = self._get_status_badge(field_node.status)
+            label_parts.append(status_badge)
+
+        label = " ".join(label_parts)
+        parent.add(label)
+
+    def _get_status_badge(self, status: NodeStatus) -> str:
+        """Get status badge text with color.
+
+        Args:
+            status: Node status
+
+        Returns:
+            Rich-formatted status badge
+        """
+        if status == NodeStatus.ADDED:
+            return f"[{self.theme.added_color}][+][/{self.theme.added_color}]"
+        elif status == NodeStatus.REMOVED:
+            return f"[{self.theme.removed_color}][-][/{self.theme.removed_color}]"
+        elif status == NodeStatus.CHANGED:
+            return f"[{self.theme.changed_color}][~][/{self.theme.changed_color}]"
+        elif status == NodeStatus.UNCHANGED:
+            return ""  # No badge for unchanged
+        else:
+            return ""  # No badge for normal

metaxy/graph/diff/rendering/theme.py

@@ -0,0 +1,48 @@
+"""Theme system for graph rendering."""
+
+from pydantic import BaseModel, Field
+
+
+class Theme(BaseModel):
+    """Color theme for graph rendering.
+
+    Unified theme for all rendering backends (terminal, Mermaid, Graphviz).
+    All colors use hex format (e.g., "#FF5733") for consistency.
+    Rich terminal supports hex via markup: [#FF0000]text[/#FF0000]
+    """
+
+    # Normal mode colors
+    feature_color: str = Field(
+        default="#00FFFF", description="Feature node color (cyan)"
+    )
+    field_color: str = Field(default="#5F87AF", description="Field color (steel blue)")
+    version_color: str = Field(
+        default="#FFFF00", description="Version info color (yellow)"
+    )
+    edge_color: str = Field(
+        default="#808080", description="Edge/dependency color (gray)"
+    )
+    snapshot_color: str = Field(
+        default="#FF00FF", description="Snapshot info color (magenta)"
+    )
+
+    # Diff mode - node/edge colors
+    added_color: str = Field(default="#00FF00", description="Added items (green)")
+    removed_color: str = Field(default="#FF0000", description="Removed items (red)")
+    changed_color: str = Field(default="#FFAA00", description="Changed items (orange)")
+    unchanged_color: str = Field(
+        default="#808080", description="Unchanged items (gray)"
+    )
+
+    # Version transition colors (for showing old→new in diffs)
+    old_version_color: str = Field(
+        default="#FF0000", description="Old version color (red)"
+    )
+    new_version_color: str = Field(
+        default="#00FF00", description="New version color (green)"
+    )
+
+    @classmethod
+    def default(cls) -> "Theme":
+        """Create default theme."""
+        return cls()

metaxy/graph/diff/traversal.py

@@ -0,0 +1,247 @@
+"""Graph traversal utilities."""
+
+from collections import deque
+
+from metaxy.graph.diff.models import GraphData, GraphNode
+from metaxy.models.types import FeatureKey
+
+
+class GraphWalker:
+    """Traverses and filters graph data structures.
+
+    Provides various traversal strategies:
+    - Topological sort (dependencies first)
+    - BFS from starting node
+    - Subgraph extraction with up/down filtering
+    """
+
+    def __init__(self, graph_data: GraphData):
+        """Initialize walker with graph data.
+
+        Args:
+            graph_data: Graph structure to traverse
+        """
+        self.graph_data = graph_data
+
+    def topological_sort(
+        self, nodes_to_include: set[str] | None = None
+    ) -> list[GraphNode]:
+        """Get nodes in topological order (dependencies first).
+
+        Uses stable alphabetical ordering when multiple nodes are at the same level.
+        This ensures deterministic output for diff comparisons.
+
+        Args:
+            nodes_to_include: Optional set of feature key strings to include.
+                            If None, includes all nodes.
+
+        Returns:
+            List of nodes sorted so dependencies appear before dependents
+        """
+        if nodes_to_include is None:
+            nodes_to_include = set(self.graph_data.nodes.keys())
+
+        visited = set()
+        result = []
+
+        def visit(key_str: str):
+            if key_str in visited or key_str not in nodes_to_include:
+                return
+            visited.add(key_str)
+
+            node = self.graph_data.nodes[key_str]
+
+            # Visit dependencies first, in sorted order for determinism
+            sorted_deps = sorted(
+                (dep_key.to_string() for dep_key in node.dependencies),
+                key=str.lower,  # Case-insensitive sort
+            )
+            for dep_key_str in sorted_deps:
+                if dep_key_str in nodes_to_include:
+                    visit(dep_key_str)
+
+            result.append(node)
+
+        # Visit all nodes in sorted order for deterministic traversal
+        for key_str in sorted(nodes_to_include, key=str.lower):
+            visit(key_str)
+
+        return result
+
+    def bfs_from(
+        self, start_key: FeatureKey, max_depth: int | None = None
+    ) -> list[GraphNode]:
+        """BFS traversal starting from a node.
+
+        Args:
+            start_key: Feature key to start from
+            max_depth: Maximum depth to traverse (None = unlimited)
+
+        Returns:
+            List of nodes in BFS order
+        """
+        start_key_str = start_key.to_string()
+        if start_key_str not in self.graph_data.nodes:
+            return []
+
+        visited = set()
+        result = []
+        queue = deque([(start_key_str, 0)])  # (key_str, depth)
+
+        while queue:
+            key_str, depth = queue.popleft()
+
+            if key_str in visited:
+                continue
+
+            if max_depth is not None and depth > max_depth:
+                continue
+
+            visited.add(key_str)
+            node = self.graph_data.nodes[key_str]
+            result.append(node)
+
+            # Add dependencies
+            for dep_key in node.dependencies:
+                dep_key_str = dep_key.to_string()
+                if dep_key_str not in visited and dep_key_str in self.graph_data.nodes:
+                    queue.append((dep_key_str, depth + 1))
+
+        return result
+
+    def extract_subgraph(
+        self,
+        focus_key: FeatureKey,
+        up: int | None = None,
+        down: int | None = None,
+    ) -> GraphData:
+        """Extract a subgraph centered on a focus node.
+
+        Args:
+            focus_key: Feature to focus on
+            up: Number of upstream levels (dependencies) to include.
+                None = all, 0 = none
+            down: Number of downstream levels (dependents) to include.
+                None = all, 0 = none
+
+        Returns:
+            New GraphData with filtered nodes and edges
+
+        Raises:
+            ValueError: If focus_key not found in graph
+        """
+        focus_key_str = focus_key.to_string()
+        if focus_key_str not in self.graph_data.nodes:
+            raise ValueError(f"Feature '{focus_key_str}' not found in graph")
+
+        # Start with focus node
+        nodes_to_include = {focus_key_str}
+
+        # Add upstream (dependencies)
+        if up != 0:
+            max_up = None if up is None or up < 0 else up
+            upstream = self._get_upstream(focus_key_str, max_levels=max_up)
+            nodes_to_include.update(upstream)
+
+        # Add downstream (dependents)
+        if down != 0:
+            max_down = None if down is None or down < 0 else down
+            downstream = self._get_downstream(focus_key_str, max_levels=max_down)
+            nodes_to_include.update(downstream)
+
+        # Filter nodes and edges
+        filtered_nodes = {
+            k: v for k, v in self.graph_data.nodes.items() if k in nodes_to_include
+        }
+
+        filtered_edges = [
+            edge
+            for edge in self.graph_data.edges
+            if edge.from_key.to_string() in nodes_to_include
+            and edge.to_key.to_string() in nodes_to_include
+        ]
+
+        return GraphData(
+            nodes=filtered_nodes,
+            edges=filtered_edges,
+            snapshot_version=self.graph_data.snapshot_version,
+            old_snapshot_version=self.graph_data.old_snapshot_version,
+        )
+
+    def _get_upstream(
+        self, start_key_str: str, max_levels: int | None = None
+    ) -> set[str]:
+        """Get upstream features (dependencies) recursively.
+
+        Args:
+            start_key_str: Feature key string to start from
+            max_levels: Maximum levels to traverse (None = unlimited)
+
+        Returns:
+            Set of upstream feature key strings
+        """
+        upstream = set()
+
+        def visit(key_str: str, level: int):
+            if key_str not in self.graph_data.nodes:
+                return
+
+            node = self.graph_data.nodes[key_str]
+
+            for dep_key in node.dependencies:
+                dep_key_str = dep_key.to_string()
+                if dep_key_str not in upstream and dep_key_str in self.graph_data.nodes:
+                    upstream.add(dep_key_str)
+                    # Only recurse if we haven't reached max level
+                    if max_levels is None or level + 1 < max_levels:
+                        visit(dep_key_str, level + 1)
+
+        visit(start_key_str, 0)
+        return upstream
+
+    def _get_downstream(
+        self, start_key_str: str, max_levels: int | None = None
+    ) -> set[str]:
+        """Get downstream features (dependents) recursively.
+
+        Args:
+            start_key_str: Feature key string to start from
+            max_levels: Maximum levels to traverse (None = unlimited)
+
+        Returns:
+            Set of downstream feature key strings
+        """
+        # Build reverse dependency map (feature -> dependents)
+        dependents_map: dict[str, list[str]] = {}
+        for node in self.graph_data.nodes.values():
+            for dep_key in node.dependencies:
+                dep_key_str = dep_key.to_string()
+                if dep_key_str not in dependents_map:
+                    dependents_map[dep_key_str] = []
+                dependents_map[dep_key_str].append(node.key.to_string())
+
+        downstream = set()
+
+        def visit(key_str: str, level: int):
+            if key_str not in dependents_map:
+                return
+
+            for dependent_key_str in dependents_map[key_str]:
+                if dependent_key_str not in downstream:
+                    downstream.add(dependent_key_str)
+                    # Only recurse if we haven't reached max level
+                    if max_levels is None or level + 1 < max_levels:
+                        visit(dependent_key_str, level + 1)
+
+        visit(start_key_str, 0)
+        return downstream
+
+    def get_root_nodes(self) -> list[GraphNode]:
+        """Get all root nodes (nodes with no dependencies).
+
+        Returns:
+            List of root nodes
+        """
+        return [
+            node for node in self.graph_data.nodes.values() if not node.dependencies
+        ]

metaxy/graph/utils.py

@@ -0,0 +1,58 @@
+"""Shared utilities for graph rendering and formatting."""
+
+from metaxy.models.types import FeatureKey, FieldKey
+
+
+def sanitize_mermaid_id(s: str) -> str:
+    """Sanitize string for use as Mermaid node ID.
+
+    Replaces characters that are invalid in Mermaid identifiers.
+
+    Args:
+        s: String to sanitize
+
+    Returns:
+        Sanitized string safe for use as Mermaid node ID
+    """
+    return s.replace("/", "_").replace("-", "_").replace("__", "_")
+
+
+def format_hash(hash_str: str, length: int = 8) -> str:
+    """Format hash string with optional truncation.
+
+    Args:
+        hash_str: Full hash string
+        length: Number of characters to show (0 for full hash)
+
+    Returns:
+        Truncated or full hash string
+    """
+    if length == 0 or length >= len(hash_str):
+        return hash_str
+    return hash_str[:length]
+
+
+def format_feature_key(key: FeatureKey) -> str:
+    """Format feature key for display.
+
+    Uses / separator for better readability.
+
+    Args:
+        key: Feature key
+
+    Returns:
+        Formatted string like "my/feature/key"
+    """
+    return "/".join(key)
+
+
+def format_field_key(key: FieldKey) -> str:
+    """Format field key for display.
+
+    Args:
+        key: Field key
+
+    Returns:
+        Formatted string like "field_name"
+    """
+    return "/".join(key)

metaxy/metadata_store/__init__.py

@@ -0,0 +1,31 @@
+"""Metadata store for feature pipeline management."""
+
+from metaxy.metadata_store.base import MetadataStore
+from metaxy.metadata_store.exceptions import (
+    DependencyError,
+    FeatureNotFoundError,
+    FieldNotFoundError,
+    HashAlgorithmNotSupportedError,
+    MetadataSchemaError,
+    MetadataStoreError,
+    StoreNotOpenError,
+)
+from metaxy.metadata_store.memory import InMemoryMetadataStore
+from metaxy.metadata_store.system_tables import (
+    FEATURE_VERSIONS_KEY,
+    allow_feature_version_override,
+)
+
+__all__ = [
+    "MetadataStore",
+    "InMemoryMetadataStore",
+    "MetadataStoreError",
+    "FeatureNotFoundError",
+    "FieldNotFoundError",
+    "MetadataSchemaError",
+    "DependencyError",
+    "StoreNotOpenError",
+    "HashAlgorithmNotSupportedError",
+    "FEATURE_VERSIONS_KEY",
+    "allow_feature_version_override",
+]

metaxy/metadata_store/_protocols.py

@@ -0,0 +1,38 @@
+"""Internal protocols for metadata store components."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, Protocol
+
+import narwhals as nw
+import polars as pl
+
+from metaxy.models.types import FeatureKey
+
+
+class MetadataStoreProtocol(Protocol):
+    """Protocol defining the interface needed by SystemTableStorage.
+
+    This protocol breaks the circular dependency between base.py and system_tables.py
+    by defining only the methods that SystemTableStorage actually uses.
+    """
+
+    def _write_metadata_impl(
+        self,
+        feature_key: FeatureKey,
+        df: pl.DataFrame,
+    ) -> None:
+        """Write metadata for a feature key."""
+        ...
+
+    def _read_metadata_native(
+        self,
+        feature: FeatureKey,
+        *,
+        feature_version: str | None = None,
+        filters: Sequence[nw.Expr] | None = None,
+        columns: Sequence[str] | None = None,
+    ) -> nw.LazyFrame[Any] | None:
+        """Read metadata from this store only (no fallback)."""
+        ...