tracepipe 0.2.0__py3-none-any.whl

tracepipe/safety.py

@@ -0,0 +1,178 @@
+# tracepipe/safety.py
+"""
+Safe instrumentation wrappers.
+
+CRITICAL: Execute original FIRST, capture lineage SECOND, return ALWAYS.
+"""
+
+import inspect
+import traceback
+import warnings
+from functools import wraps
+from typing import Callable
+
+from .context import get_context
+
+
+class TracePipeWarning(UserWarning):
+    """Warning for non-fatal instrumentation issues."""
+
+    pass
+
+
+class TracePipeError(Exception):
+    """Error raised in strict mode."""
+
+    pass
+
+
+def get_caller_info(skip_frames: int = 2) -> tuple:
+    """
+    Get caller's file and line number, skipping library frames.
+
+    Walks up the stack to find the first frame that's in user code,
+    not in pandas, numpy, or tracepipe internals.
+
+    Args:
+        skip_frames: Minimum frames to skip (default 2 for wrapper + this func)
+
+    Returns:
+        (filename, line_number) or (None, None)
+    """
+    import os
+
+    # Get tracepipe package directory to skip it specifically
+    tracepipe_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    tracepipe_pkg = os.path.join(tracepipe_dir, "tracepipe")
+
+    # Paths/patterns to skip (library code)
+    SKIP_PATTERNS = (
+        "/pandas/",
+        "/numpy/",
+        "site-packages",
+        "<frozen",
+        "<string>",
+    )
+
+    try:
+        frame = inspect.currentframe()
+        # Skip minimum frames first
+        for _ in range(skip_frames + 1):
+            if frame is None:
+                return None, None
+            frame = frame.f_back
+
+        # Now walk up until we find user code
+        max_depth = 20  # Safety limit
+        for _ in range(max_depth):
+            if frame is None:
+                return None, None
+
+            filename = frame.f_code.co_filename
+            abs_filename = os.path.abspath(filename)
+
+            # Check if this is tracepipe package code
+            if abs_filename.startswith(tracepipe_pkg):
+                frame = frame.f_back
+                continue
+
+            # Check if this is other library code
+            is_library = any(pattern in filename for pattern in SKIP_PATTERNS)
+
+            if not is_library:
+                return filename, frame.f_lineno
+
+            frame = frame.f_back
+
+        return None, None
+    except Exception:
+        return None, None
+    finally:
+        del frame
+
+
+def _make_wrapper(
+    method_name: str, original_method: Callable, capture_func: Callable, mode: str = "standard"
+) -> Callable:
+    """
+    Factory for pandas method wrappers with lineage capture.
+
+    CRITICAL: Execute original FIRST, capture lineage SECOND, return ALWAYS.
+
+    Args:
+        method_name: Name for error messages
+        original_method: The original pandas method
+        capture_func: func(self, args, kwargs, result, ctx, method_name)
+        mode: "standard", "filter", or "inplace"
+    """
+
+    @wraps(original_method)
+    def wrapper(self, *args, **kwargs):
+        ctx = get_context()
+
+        # === PRE-EXECUTION SETUP ===
+        before_snapshot = None
+
+        if mode == "filter" and ctx.enabled:
+            ctx._filter_op_depth += 1
+        elif mode == "inplace" and ctx.enabled and kwargs.get("inplace", False):
+            try:
+                before_snapshot = self.copy()
+            except Exception:
+                pass
+
+        # === EXECUTE ORIGINAL (SACRED) ===
+        try:
+            result = original_method(self, *args, **kwargs)
+        finally:
+            if mode == "filter" and ctx.enabled:
+                ctx._filter_op_depth -= 1
+
+        # === CAPTURE LINEAGE (SIDE EFFECT) ===
+        if ctx.enabled:
+            try:
+                if mode == "inplace" and kwargs.get("inplace", False):
+                    if before_snapshot is not None:
+                        capture_func(before_snapshot, args, kwargs, self, ctx, method_name)
+                elif mode == "inplace" and result is not None:
+                    capture_func(self, args, kwargs, result, ctx, method_name)
+                else:
+                    capture_func(self, args, kwargs, result, ctx, method_name)
+            except Exception as e:
+                if ctx.config.strict_mode:
+                    raise TracePipeError(
+                        f"Instrumentation failed for {method_name}: {e}\n"
+                        f"{traceback.format_exc()}"
+                    ) from e
+                else:
+                    warnings.warn(
+                        f"TracePipe: {method_name} instrumentation failed: {e}. "
+                        f"Lineage may be incomplete.",
+                        TracePipeWarning,
+                    )
+
+        # === RETURN RESULT (ALWAYS) ===
+        return result
+
+    return wrapper
+
+
+def wrap_pandas_method(
+    method_name: str, original_method: Callable, capture_func: Callable
+) -> Callable:
+    """Wrap a pandas method with lineage capture."""
+    return _make_wrapper(method_name, original_method, capture_func, mode="standard")
+
+
+def wrap_pandas_filter_method(
+    method_name: str, original_method: Callable, capture_func: Callable
+) -> Callable:
+    """Wrap a pandas filter method (dropna, drop_duplicates, etc.)."""
+    return _make_wrapper(method_name, original_method, capture_func, mode="filter")
+
+
+def wrap_pandas_method_inplace(
+    method_name: str, original_method: Callable, capture_func: Callable
+) -> Callable:
+    """Wrap a pandas method that supports inplace=True."""
+    return _make_wrapper(method_name, original_method, capture_func, mode="inplace")

tracepipe/storage/__init__.py

@@ -0,0 +1,13 @@
+# tracepipe/storage/__init__.py
+"""Storage backends and row identity strategies."""
+
+from .base import LineageBackend, RowIdentityStrategy
+from .lineage_store import InMemoryLineageStore
+from .row_identity import PandasRowIdentity
+
+__all__ = [
+    "LineageBackend",
+    "RowIdentityStrategy",
+    "InMemoryLineageStore",
+    "PandasRowIdentity",
+]

tracepipe/storage/base.py

@@ -0,0 +1,174 @@
+# tracepipe/storage/base.py
+"""
+Protocol definitions for TracePipe storage backends.
+
+These protocols enable:
+- Swappable storage backends (InMemory, SQLite, Delta Lake)
+- Engine-specific row identity strategies (Pandas, Polars, Spark)
+- Easy testing with mock implementations
+
+To add a new backend, implement LineageBackend.
+To support a new DataFrame engine, implement RowIdentityStrategy.
+"""
+
+from typing import Any, Optional, Protocol, runtime_checkable
+
+from ..core import ChangeType, CompletenessLevel, LineageGaps, TracePipeConfig
+
+
+@runtime_checkable
+class LineageBackend(Protocol):
+    """
+    Protocol for lineage storage backends.
+
+    Implementations:
+    - InMemoryLineageStore (default, v0.2.0)
+    - SQLiteBackend (future)
+    - DeltaLakeBackend (future)
+    """
+
+    config: TracePipeConfig
+
+    def append_diff(
+        self,
+        step_id: int,
+        row_id: int,
+        col: str,
+        old_val: Any,
+        new_val: Any,
+        change_type: ChangeType,
+    ) -> None:
+        """Append a single cell diff."""
+        ...
+
+    def append_diff_batch(
+        self, step_id: int, diffs: list[tuple], check_threshold: bool = True
+    ) -> int:
+        """Batch append diffs. Returns count appended."""
+        ...
+
+    def append_step(
+        self,
+        operation: str,
+        stage: Optional[str],
+        code_file: Optional[str],
+        code_line: Optional[int],
+        params: dict[str, Any],
+        input_shape: Optional[tuple],
+        output_shape: Optional[tuple],
+        completeness: CompletenessLevel = CompletenessLevel.FULL,
+        is_mass_update: bool = False,
+        rows_affected: int = 0,
+    ) -> int:
+        """Append step metadata. Returns step_id."""
+        ...
+
+    def append_aggregation(
+        self,
+        step_id: int,
+        group_column: str,
+        membership: dict[str, list[int]],
+        agg_functions: dict[str, str],
+    ) -> None:
+        """Record aggregation group membership."""
+        ...
+
+    def get_row_history(self, row_id: int) -> list[dict]:
+        """Get all events for a specific row."""
+        ...
+
+    def get_dropped_rows(self, step_id: Optional[int] = None) -> list[int]:
+        """Get dropped row IDs, optionally filtered by step."""
+        ...
+
+    def get_dropped_by_step(self) -> dict[str, int]:
+        """Get count of dropped rows per operation."""
+        ...
+
+    def get_group_members(self, group_key: str) -> Optional[dict]:
+        """Get rows that contributed to a group."""
+        ...
+
+    def compute_gaps(self, row_id: int) -> LineageGaps:
+        """Compute lineage gaps for a row."""
+        ...
+
+    def should_track_cell_diffs(self, affected_count: int) -> bool:
+        """Return False for mass updates exceeding threshold."""
+        ...
+
+    def to_json(self) -> str:
+        """Export all data as JSON string."""
+        ...
+
+    @property
+    def steps(self) -> list:
+        """Access step metadata list."""
+        ...
+
+    @property
+    def total_diff_count(self) -> int:
+        """Total diffs including spilled."""
+        ...
+
+    @property
+    def diff_count(self) -> int:
+        """In-memory diff count."""
+        ...
+
+
+@runtime_checkable
+class RowIdentityStrategy(Protocol):
+    """
+    Protocol for row identity tracking.
+
+    Implementations:
+    - PandasRowIdentity (default, v0.2.0)
+    - PolarsRowIdentity (future)
+    - SparkRowIdentity (future)
+    """
+
+    config: TracePipeConfig
+
+    def register(
+        self, df: Any, row_ids: Optional[Any] = None, warn_duplicate_index: bool = True
+    ) -> Any:
+        """Register a DataFrame and assign/return row IDs."""
+        ...
+
+    def get_ids(self, df: Any) -> Optional[Any]:
+        """Get row IDs for a DataFrame, or None if not tracked."""
+        ...
+
+    def propagate(self, source_df: Any, result_df: Any) -> Optional[Any]:
+        """Propagate row IDs from source to result DataFrame."""
+        ...
+
+    def get_dropped_ids(self, source_df: Any, result_df: Any) -> set:
+        """Get row IDs that were dropped between source and result."""
+        ...
+
+    def strip_hidden_column(self, df: Any) -> Any:
+        """Remove hidden column for export."""
+        ...
+
+    def cleanup(self) -> None:
+        """Remove stale entries."""
+        ...
+
+
+# === FACTORY FUNCTIONS ===
+
+
+def create_default_backend(config: TracePipeConfig) -> "LineageBackend":
+    """Create the default in-memory backend."""
+    from .lineage_store import InMemoryLineageStore
+
+    return InMemoryLineageStore(config)
+
+
+def create_default_identity(config: TracePipeConfig) -> "RowIdentityStrategy":
+    """Create the default pandas row identity strategy."""
+    from .row_identity import PandasRowIdentity
+
+    return PandasRowIdentity(config)