tracepipe 0.2.0__py3-none-any.whl

tracepipe/context.py

@@ -0,0 +1,98 @@
+# tracepipe/context.py
+"""
+Thread-safe context for TracePipe state.
+
+Each thread gets its own context via threading.local().
+"""
+
+import threading
+from typing import Optional
+
+from .core import TracePipeConfig
+from .storage.base import (
+    LineageBackend,
+    RowIdentityStrategy,
+    create_default_backend,
+    create_default_identity,
+)
+
+# Thread-local storage for context
+_thread_local = threading.local()
+
+
+class TracePipeContext:
+    """
+    Per-thread context for TracePipe state.
+
+    Thread Safety:
+    - Each thread gets its own context via threading.local()
+    - Shared state (if needed) must use locks
+    - This design supports concurrent notebook cells but NOT
+      parallel pandas operations on shared DataFrames
+
+    Extensibility:
+    - Pass custom `backend` for alternative storage (SQLite, Delta Lake)
+    - Pass custom `identity` for alternative engines (Polars, Spark)
+    """
+
+    def __init__(
+        self,
+        config: Optional[TracePipeConfig] = None,
+        backend: Optional[LineageBackend] = None,
+        identity: Optional[RowIdentityStrategy] = None,
+    ):
+        self.config = config or TracePipeConfig()
+        self.enabled: bool = False
+
+        # Use provided backends or create defaults
+        self.store: LineageBackend = backend or create_default_backend(self.config)
+        self.row_manager: RowIdentityStrategy = identity or create_default_identity(self.config)
+
+        self.watched_columns: set[str] = set()
+        self.current_stage: Optional[str] = None
+
+        # Nested filter operation tracking (prevents double-counting drops)
+        # When > 0, __getitem__[mask] skips capture (parent op will capture)
+        self._filter_op_depth: int = 0
+
+        # GroupBy state stack (supports nesting)
+        self._groupby_stack: list[dict] = []
+
+    def push_groupby(self, state: dict) -> None:
+        """Push groupby state for nested operations."""
+        self._groupby_stack.append(state)
+
+    def pop_groupby(self) -> Optional[dict]:
+        """Pop most recent groupby state."""
+        return self._groupby_stack.pop() if self._groupby_stack else None
+
+    def peek_groupby(self) -> Optional[dict]:
+        """Peek at current groupby state without removing."""
+        return self._groupby_stack[-1] if self._groupby_stack else None
+
+    def clear_groupby_for_source(self, source_id: int) -> None:
+        """
+        Clear any groupby state for a given source DataFrame.
+
+        Called when a new groupby() is performed on the same DataFrame,
+        which invalidates any previous groupby state for that source.
+        """
+        self._groupby_stack = [s for s in self._groupby_stack if s.get("source_id") != source_id]
+
+
+def get_context() -> TracePipeContext:
+    """Get the current thread's TracePipe context."""
+    if not hasattr(_thread_local, "context"):
+        _thread_local.context = TracePipeContext()
+    return _thread_local.context
+
+
+def set_context(ctx: TracePipeContext) -> None:
+    """Set context for current thread (used in testing)."""
+    _thread_local.context = ctx
+
+
+def reset_context() -> None:
+    """Reset context for current thread."""
+    if hasattr(_thread_local, "context"):
+        del _thread_local.context

tracepipe/core.py

@@ -0,0 +1,122 @@
+# tracepipe/core.py
+"""
+Core types, enums, and configuration for TracePipe.
+"""
+
+import os
+from dataclasses import dataclass, field
+from enum import IntEnum
+from typing import Any, Optional
+
+
+class ChangeType(IntEnum):
+    """Types of changes tracked by TracePipe."""
+
+    MODIFIED = 0
+    DROPPED = 1
+    ADDED = 2
+    REORDERED = 3
+
+
+class CompletenessLevel(IntEnum):
+    """
+    Indicates how completely an operation's internals are tracked.
+
+    FULL: Completely tracked (e.g., fillna, dropna)
+    PARTIAL: Output tracked, internals unknown (e.g., apply, pipe)
+    UNKNOWN: Lineage reset (e.g., merge, concat)
+    """
+
+    FULL = 0
+    PARTIAL = 1
+    UNKNOWN = 2
+
+
+@dataclass
+class TracePipeConfig:
+    """Configuration with sensible defaults."""
+
+    max_diffs_in_memory: int = 500_000
+    max_diffs_per_step: int = 100_000
+    max_group_membership_size: int = 100_000  # Store count-only above this threshold
+    strict_mode: bool = False
+    auto_watch: bool = False
+    auto_watch_null_threshold: float = 0.01
+    spillover_dir: str = ".tracepipe"
+    use_hidden_column: bool = False
+    warn_on_duplicate_index: bool = True
+    cleanup_spillover_on_disable: bool = True
+
+    @classmethod
+    def from_env(cls) -> "TracePipeConfig":
+        """Create config from environment variables."""
+        return cls(
+            max_diffs_in_memory=int(os.environ.get("TRACEPIPE_MAX_DIFFS", 500_000)),
+            max_diffs_per_step=int(os.environ.get("TRACEPIPE_MAX_DIFFS_PER_STEP", 100_000)),
+            strict_mode=os.environ.get("TRACEPIPE_STRICT", "0") == "1",
+            auto_watch=os.environ.get("TRACEPIPE_AUTO_WATCH", "0") == "1",
+            use_hidden_column=os.environ.get("TRACEPIPE_HIDDEN_COL", "0") == "1",
+        )
+
+
+@dataclass
+class StepMetadata:
+    """Metadata for a single pipeline step."""
+
+    step_id: int
+    operation: str
+    stage: Optional[str]
+    timestamp: float
+    code_file: Optional[str]
+    code_line: Optional[int]
+    params: dict[str, Any]
+    input_shape: Optional[tuple]
+    output_shape: Optional[tuple]
+    is_mass_update: bool = False
+    rows_affected: int = 0
+    completeness: CompletenessLevel = CompletenessLevel.FULL
+
+
+@dataclass
+class AggregationMapping:
+    """Tracks which rows contributed to an aggregation group."""
+
+    step_id: int
+    group_column: str
+    membership: dict[str, list[int]]  # {group_key: [row_ids]}
+    agg_functions: dict[str, str]
+
+
+@dataclass
+class LineageGap:
+    """Represents a gap in lineage tracking."""
+
+    step_id: int
+    operation: str
+    reason: str
+
+
+@dataclass
+class LineageGaps:
+    """Collection of lineage gaps for a row."""
+
+    gaps: list[LineageGap] = field(default_factory=list)
+
+    @property
+    def has_gaps(self) -> bool:
+        """Return True if there are any gaps."""
+        return len(self.gaps) > 0
+
+    @property
+    def is_fully_tracked(self) -> bool:
+        """Return True if lineage is complete."""
+        return len(self.gaps) == 0
+
+    def summary(self) -> str:
+        """Return a human-readable summary."""
+        if self.is_fully_tracked:
+            return "Fully tracked"
+        elif len(self.gaps) == 1:
+            return f"1 step has limited visibility: {self.gaps[0].operation}"
+        else:
+            return f"{len(self.gaps)} steps have limited visibility"

tracepipe/instrumentation/__init__.py

@@ -0,0 +1,6 @@
+# tracepipe/instrumentation/__init__.py
+"""Instrumentation for various data processing libraries."""
+
+from .pandas_inst import instrument_pandas, uninstrument_pandas
+
+__all__ = ["instrument_pandas", "uninstrument_pandas"]