tracepipe 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

tracepipe/instrumentation/pandas_inst.py

@@ -1,6 +1,11 @@
 # tracepipe/instrumentation/pandas_inst.py
 """
 Pandas DataFrame instrumentation for row-level lineage tracking.
+
+This module wires the filter and merge capture modules to pandas methods.
+The actual capture logic is in:
+- filter_capture.py: Mask-first filter capture with FULL/PARTIAL completeness
+- merge_capture.py: Position column injection for merge provenance
 """
 
 import warnings
@@ -15,56 +20,26 @@ from ..core import ChangeType, CompletenessLevel
 from ..safety import (
     TracePipeWarning,
     get_caller_info,
-    wrap_pandas_filter_method,
     wrap_pandas_method,
     wrap_pandas_method_inplace,
 )
 from ..utils.value_capture import find_changed_indices_vectorized
+from .apply_capture import instrument_apply_pipe, uninstrument_apply_pipe
+
+# Import new capture modules
+from .filter_capture import wrap_filter_method, wrap_getitem_filter
+from .indexer_capture import instrument_indexers, uninstrument_indexers
+from .merge_capture import (
+    wrap_concat_with_lineage,
+    wrap_join_with_lineage,
+    wrap_merge_with_lineage,
+)
+from .series_capture import instrument_series, uninstrument_series
 
 # Store original methods for restore
 _originals: dict[str, Any] = {}
 
 
-# === FILTER CAPTURE ===
-
-
-def _capture_filter(
-    self: pd.DataFrame, args, kwargs, result, ctx: TracePipeContext, method_name: str
-):
-    """Capture lineage for filter operations (dropna, query, head, etc.)."""
-    if not isinstance(result, pd.DataFrame):
-        return
-
-    source_ids = ctx.row_manager.get_ids(self)
-    if source_ids is None:
-        # Auto-register if not tracked
-        ctx.row_manager.register(self)
-        source_ids = ctx.row_manager.get_ids(self)
-        if source_ids is None:
-            return
-
-    # Propagate IDs to result
-    ctx.row_manager.propagate(self, result)
-
-    # Find dropped rows (returns numpy array for performance)
-    dropped_ids = ctx.row_manager.get_dropped_ids(self, result)
-
-    if len(dropped_ids) > 0:
-        code_file, code_line = get_caller_info(skip_frames=2)
-        step_id = ctx.store.append_step(
-            operation=f"DataFrame.{method_name}",
-            stage=ctx.current_stage,
-            code_file=code_file,
-            code_line=code_line,
-            params=_safe_params(kwargs),
-            input_shape=self.shape,
-            output_shape=result.shape,
-        )
-
-        # Bulk record all drops at once (10-50x faster than loop)
-        ctx.store.append_bulk_drops(step_id, dropped_ids)
-
-
 # === TRANSFORM CAPTURE ===
 
 
@@ -299,56 +274,6 @@ def _capture_agg(self, args, kwargs, result, ctx: TracePipeContext, method_name:
         ctx.row_manager.register(result)
 
 
-# === MERGE/CONCAT (UNKNOWN - OUT OF SCOPE) ===
-
-
-def _capture_merge(
-    self: pd.DataFrame, args, kwargs, result, ctx: TracePipeContext, method_name: str
-):
-    """Mark merge as UNKNOWN completeness and reset lineage."""
-    code_file, code_line = get_caller_info(skip_frames=2)
-    ctx.store.append_step(
-        operation=f"DataFrame.{method_name}",
-        stage=ctx.current_stage,
-        code_file=code_file,
-        code_line=code_line,
-        params={"how": kwargs.get("how", "inner")},
-        input_shape=self.shape,
-        output_shape=result.shape if hasattr(result, "shape") else None,
-        completeness=CompletenessLevel.UNKNOWN,
-    )
-
-    warnings.warn(
-        f"TracePipe: {method_name}() resets row lineage. "
-        f"Rows in result cannot be traced back to source rows.",
-        TracePipeWarning,
-    )
-
-    # Register result with NEW IDs
-    if isinstance(result, pd.DataFrame):
-        ctx.row_manager.register(result)
-
-
-def _capture_concat(args, kwargs, result, ctx: TracePipeContext):
-    """Capture pd.concat (module-level function)."""
-    code_file, code_line = get_caller_info(skip_frames=2)
-    ctx.store.append_step(
-        operation="pd.concat",
-        stage=ctx.current_stage,
-        code_file=code_file,
-        code_line=code_line,
-        params={"axis": kwargs.get("axis", 0)},
-        input_shape=None,
-        output_shape=result.shape if hasattr(result, "shape") else None,
-        completeness=CompletenessLevel.UNKNOWN,
-    )
-
-    warnings.warn("TracePipe: pd.concat() resets row lineage.", TracePipeWarning)
-
-    if isinstance(result, pd.DataFrame):
-        ctx.row_manager.register(result)
-
-
 # === INDEX OPERATIONS ===
 
 
@@ -408,21 +333,21 @@ def _capture_sort_values(
     # Record reorder for each row
     result_ids = ctx.row_manager.get_ids(result)
     if result_ids is not None:
+        # Pre-compute position lookup (O(n) instead of O(n²))
+        source_idx_list = list(source_ids.index) if hasattr(source_ids, "index") else []
+        source_pos_map = {idx: pos for pos, idx in enumerate(source_idx_list)}
+
         for new_pos, (idx, row_id) in enumerate(result_ids.items()):
-            # Find old position
-            try:
-                old_pos = list(source_ids.index).index(idx)
-                if old_pos != new_pos:
-                    ctx.store.append_diff(
-                        step_id=step_id,
-                        row_id=int(row_id),
-                        col="__position__",
-                        old_val=old_pos,
-                        new_val=new_pos,
-                        change_type=ChangeType.REORDERED,
-                    )
-            except (ValueError, KeyError):
-                pass
+            old_pos = source_pos_map.get(idx)
+            if old_pos is not None and old_pos != new_pos:
+                ctx.store.append_diff(
+                    step_id=step_id,
+                    row_id=int(row_id),
+                    col="__position__",
+                    old_val=old_pos,
+                    new_val=new_pos,
+                    change_type=ChangeType.REORDERED,
+                )
 
 
 # === COPY CAPTURE ===
@@ -458,7 +383,11 @@ def _capture_drop(
 
     - Row drops (axis=0): Track as filter operation
     - Column drops (axis=1): Track as schema change (step metadata only)
+    - Handles inplace=True (result is self, passed by wrapper)
     """
+    # Handle inplace: result is passed as self by the inplace wrapper
+    if result is None:
+        return
     if not isinstance(result, pd.DataFrame):
         return
 
@@ -509,60 +438,6 @@ def _capture_drop(
         )
 
 
-# === __getitem__ DISPATCH ===
-
-
-def _capture_getitem(
-    self: pd.DataFrame, args, kwargs, result, ctx: TracePipeContext, method_name: str
-):
-    """
-    Dispatch __getitem__ based on key type.
-
-    - df['col'] -> Series (ignore)
-    - df[['a','b']] -> DataFrame column select (propagate)
-    - df[mask] -> DataFrame row filter (track drops)
-    - df[slice] -> DataFrame row slice (track drops)
-    """
-    if len(args) != 1:
-        return
-
-    key = args[0]
-
-    # Series result - column access, not row filter
-    if isinstance(result, pd.Series):
-        return
-
-    if not isinstance(result, pd.DataFrame):
-        return
-
-    # Boolean mask - row filter
-    if isinstance(key, (pd.Series, np.ndarray)) and getattr(key, "dtype", None) is np.dtype("bool"):
-        # Skip if we're inside a named filter op (e.g., drop_duplicates)
-        # to avoid double-counting drops
-        if ctx._filter_op_depth > 0:
-            ctx.row_manager.propagate(self, result)
-            return
-        _capture_filter(self, args, kwargs, result, ctx, "__getitem__[mask]")
-        return
-
-    # List of columns - column selection
-    if isinstance(key, list):
-        ctx.row_manager.propagate(self, result)
-        return
-
-    # Slice - row selection
-    if isinstance(key, slice):
-        # Skip if we're inside a named filter op
-        if ctx._filter_op_depth > 0:
-            ctx.row_manager.propagate(self, result)
-            return
-        _capture_filter(self, args, kwargs, result, ctx, "__getitem__[slice]")
-        return
-
-    # Default: propagate
-    ctx.row_manager.propagate(self, result)
-
-
 # === __setitem__ CAPTURE ===
 
 
@@ -751,7 +626,8 @@ def _wrap_dataframe_init(original):
         original(self, *args, **kwargs)
 
         ctx = get_context()
-        if ctx.enabled:
+        # Skip registration when inside filter/export operation (prevents re-adding hidden column)
+        if ctx.enabled and ctx._filter_op_depth == 0:
             if ctx.row_manager.get_ids(self) is None:
                 ctx.row_manager.register(self)
 
@@ -768,7 +644,12 @@ def _make_export_wrapper(original):
     def wrapper(self, *args, **kwargs):
         ctx = get_context()
         if ctx.enabled:
-            clean_df = ctx.row_manager.strip_hidden_column(self)
+            # Increment filter depth to prevent tracking during column strip
+            ctx._filter_op_depth += 1
+            try:
+                clean_df = ctx.row_manager.strip_hidden_column(self)
+            finally:
+                ctx._filter_op_depth -= 1
             return original(clean_df, *args, **kwargs)
         return original(self, *args, **kwargs)
 
@@ -880,15 +761,14 @@ def instrument_pandas():
         # Already instrumented
         return
 
-    # === DataFrame filter methods ===
-    # Use wrap_pandas_filter_method to prevent double-counting when
-    # methods like drop_duplicates internally call __getitem__
+    # === DataFrame filter methods (using new mask-first capture) ===
+    # wrap_filter_method provides FULL/PARTIAL completeness tracking
     filter_methods = ["dropna", "drop_duplicates", "query", "head", "tail", "sample"]
     for method_name in filter_methods:
         if hasattr(pd.DataFrame, method_name):
             original = getattr(pd.DataFrame, method_name)
             _originals[f"DataFrame.{method_name}"] = original
-            wrapped = wrap_pandas_filter_method(method_name, original, _capture_filter)
+            wrapped = wrap_filter_method(method_name, original)
             setattr(pd.DataFrame, method_name, wrapped)
 
     # === DataFrame transform methods (with inplace support) ===
@@ -908,9 +788,9 @@ def instrument_pandas():
     _originals["DataFrame.copy"] = pd.DataFrame.copy
     pd.DataFrame.copy = wrap_pandas_method("copy", pd.DataFrame.copy, _capture_copy)
 
-    # === drop (row/column removal) ===
+    # === drop (row/column removal, supports inplace=True) ===
     _originals["DataFrame.drop"] = pd.DataFrame.drop
-    pd.DataFrame.drop = wrap_pandas_method("drop", pd.DataFrame.drop, _capture_drop)
+    pd.DataFrame.drop = wrap_pandas_method_inplace("drop", pd.DataFrame.drop, _capture_drop)
 
     # === apply/pipe ===
     _originals["DataFrame.apply"] = pd.DataFrame.apply
@@ -933,12 +813,14 @@ def instrument_pandas():
             wrapped = wrap_pandas_method(agg_method, original, _capture_agg)
             setattr(DataFrameGroupBy, agg_method, wrapped)
 
-    # === merge ===
+    # === merge (using new position column injection capture) ===
+    # wrap_merge_with_lineage provides full provenance in DEBUG mode
     _originals["DataFrame.merge"] = pd.DataFrame.merge
-    pd.DataFrame.merge = wrap_pandas_method("merge", pd.DataFrame.merge, _capture_merge)
+    pd.DataFrame.merge = wrap_merge_with_lineage(pd.DataFrame.merge)
 
+    # === join (using new join wrapper) ===
     _originals["DataFrame.join"] = pd.DataFrame.join
-    pd.DataFrame.join = wrap_pandas_method("join", pd.DataFrame.join, _capture_merge)
+    pd.DataFrame.join = wrap_join_with_lineage(pd.DataFrame.join)
 
     # === Index operations ===
     _originals["DataFrame.reset_index"] = pd.DataFrame.reset_index
@@ -956,11 +838,9 @@ def instrument_pandas():
         "sort_values", pd.DataFrame.sort_values, _capture_sort_values
     )
 
-    # === __getitem__ ===
+    # === __getitem__ (using new filter capture for boolean indexing) ===
     _originals["DataFrame.__getitem__"] = pd.DataFrame.__getitem__
-    pd.DataFrame.__getitem__ = wrap_pandas_method(
-        "__getitem__", pd.DataFrame.__getitem__, _capture_getitem
-    )
+    pd.DataFrame.__getitem__ = wrap_getitem_filter(pd.DataFrame.__getitem__)
 
     # === __setitem__ (column assignment) ===
     _originals["DataFrame.__setitem__"] = pd.DataFrame.__setitem__
@@ -993,23 +873,26 @@ def instrument_pandas():
     _originals["DataFrame.to_parquet"] = pd.DataFrame.to_parquet
     pd.DataFrame.to_parquet = _make_export_wrapper(pd.DataFrame.to_parquet)
 
-    # === pd.concat ===
+    # === pd.concat (using new concat wrapper) ===
     _originals["pd.concat"] = pd.concat
+    pd.concat = wrap_concat_with_lineage(_originals["pd.concat"])
 
-    def wrapped_concat(*args, **kwargs):
-        result = _originals["pd.concat"](*args, **kwargs)
-        ctx = get_context()
-        if ctx.enabled:
-            _capture_concat(args, kwargs, result, ctx)
-        return result
-
-    pd.concat = wrapped_concat
+    # === Phase 6: Extended operation support ===
+    # Note: These modules handle their own original storage
+    instrument_indexers()
+    instrument_series()
+    instrument_apply_pipe()
 
 
 def uninstrument_pandas():
     """Restore original pandas methods."""
     global _originals
 
+    # Uninstrument Phase 6 modules first (reverse order)
+    uninstrument_apply_pipe()
+    uninstrument_series()
+    uninstrument_indexers()
+
     for key, original in _originals.items():
         parts = key.split(".")
         if parts[0] == "pd":

tracepipe/instrumentation/series_capture.py

@@ -0,0 +1,331 @@
+# tracepipe/instrumentation/series_capture.py
+"""
+Series method instrumentation for TracePipe.
+
+Challenge: Series operations are often chained and may not be assigned back.
+Solution: Track when Series is extracted, wrap common methods, capture on assignment.
+
+Operations tracked:
+| Pattern                           | Tracking                    | Completeness |
+|-----------------------------------|-----------------------------|--------------|
+| df['col'].fillna(val)             | Method call + assignment    | FULL         |
+| df['col'].replace(...)            | Method call + assignment    | FULL         |
+| df['col'].str.upper()             | Method call + assignment    | FULL         |
+| df['col'].dt.year                 | Method call + assignment    | FULL         |
+| df['col'].apply(func)             | Before/after diff           | PARTIAL      |
+| df['col'] = series                | Assignment diff             | FULL         |
+
+Key insight: We track at ASSIGNMENT time, not method call time.
+This handles arbitrary chains: df['col'] = df['other'].str.strip().str.upper()
+"""
+
+import warnings
+import weakref
+from functools import wraps
+
+import pandas as pd
+
+from ..context import get_context
+from ..core import ChangeType, CompletenessLevel
+from ..safety import TracePipeWarning, get_caller_info
+
+
+class TrackedSeries(pd.Series):
+    """
+    Series subclass that tracks its origin DataFrame and column.
+
+    When assigned back to a DataFrame, we can compute the diff.
+
+    Note: This is created only when extracting from a tracked DataFrame.
+    Regular Series operations remain unchanged.
+
+    Memory Safety:
+    - _tp_source_df_ref is a weakref to prevent memory leaks
+    - Source DataFrame can be garbage collected independently
+    """
+
+    _metadata = ["_tp_source_df_ref", "_tp_source_col", "_tp_source_rids", "_tp_last_op"]
+
+    @property
+    def _constructor(self):
+        return TrackedSeries
+
+    @property
+    def _constructor_expanddim(self):
+        return pd.DataFrame
+
+    @property
+    def _tp_source_df(self):
+        """Get source DataFrame from weakref (may return None if GC'd)."""
+        ref = getattr(self, "_tp_source_df_ref", None)
+        if ref is not None:
+            return ref()
+        return None
+
+    @_tp_source_df.setter
+    def _tp_source_df(self, df):
+        """Store source DataFrame as weakref."""
+        if df is not None:
+            self._tp_source_df_ref = weakref.ref(df)
+        else:
+            self._tp_source_df_ref = None
+
+
+def wrap_series_extraction():
+    """
+    Wrap DataFrame.__getitem__ to return TrackedSeries for single column access.
+
+    This allows us to track the origin of Series that may be modified and assigned back.
+    """
+    original_getitem = pd.DataFrame.__getitem__
+
+    @wraps(original_getitem)
+    def tracked_getitem(self, key):
+        result = original_getitem(self, key)
+
+        ctx = get_context()
+        if not ctx.enabled:
+            return result
+
+        # Skip internal tracking operations to avoid recursion
+        if ctx._filter_op_depth > 0:
+            return result
+
+        # Skip internal tracepipe columns
+        if isinstance(key, str) and key.startswith("__tracepipe"):
+            return result
+
+        # Only wrap single-column Series access
+        if isinstance(key, str) and isinstance(result, pd.Series):
+            rids = ctx.row_manager.get_ids_array(self)
+            if rids is not None:
+                # Convert to TrackedSeries
+                tracked = TrackedSeries(result)
+                tracked._tp_source_df = self
+                tracked._tp_source_col = key
+                tracked._tp_source_rids = rids.copy()
+                return tracked
+
+        return result
+
+    pd.DataFrame.__getitem__ = tracked_getitem
+    pd.DataFrame._tp_original_getitem_series = original_getitem
+
+
+def wrap_series_assignment():
+    """
+    Wrap DataFrame.__setitem__ to capture diffs when assigning Series.
+
+    Handles:
+    - df['col'] = series  (where series may have been modified)
+    - df['col'] = scalar  (broadcast assignment)
+    - df['col'] = array   (direct assignment)
+    """
+    original_setitem = pd.DataFrame.__setitem__
+
+    @wraps(original_setitem)
+    def tracked_setitem(self, key, value):
+        ctx = get_context()
+
+        # Capture before state for watched columns
+        before_values = None
+        if (
+            ctx.enabled
+            and isinstance(key, str)
+            and key in ctx.watched_columns
+            and key in self.columns
+        ):
+            rids = ctx.row_manager.get_ids_array(self)
+            if rids is not None:
+                before_values = {
+                    "rids": rids.copy(),
+                    "values": self[key].values.copy(),
+                }
+
+        # Always run original
+        original_setitem(self, key, value)
+
+        if not ctx.enabled:
+            return
+
+        if before_values is None:
+            return
+
+        try:
+            _capture_series_assignment(self, key, value, before_values, ctx)
+        except Exception as e:
+            if ctx.config.strict_mode:
+                raise
+            warnings.warn(f"TracePipe: Series assignment capture failed: {e}", TracePipeWarning)
+
+    pd.DataFrame.__setitem__ = tracked_setitem
+    pd.DataFrame._tp_original_setitem_series = original_setitem
+
+
+def _capture_series_assignment(df, key, value, before_values, ctx):
+    """Capture diffs from Series assignment."""
+    from ..utils.value_capture import values_equal
+
+    store = ctx.store
+    rids = before_values["rids"]
+    old_vals = before_values["values"]
+    new_vals = df[key].values
+
+    # Determine completeness based on value type
+    if isinstance(value, TrackedSeries):
+        # Can trace back to source
+        completeness = CompletenessLevel.FULL
+        operation = f"Series.{_infer_series_operation(value)}"
+    elif hasattr(value, "apply") or callable(value):
+        completeness = CompletenessLevel.PARTIAL
+        operation = "Series.transform"
+    else:
+        completeness = CompletenessLevel.FULL
+        operation = "DataFrame[]="
+
+    code_file, code_line = get_caller_info(skip_frames=4)
+    step_id = store.append_step(
+        operation=operation,
+        stage=ctx.current_stage,
+        code_file=code_file,
+        code_line=code_line,
+        params={"column": key},
+        input_shape=df.shape,
+        output_shape=df.shape,
+        completeness=completeness,
+    )
+
+    # Track diffs for changed values
+    for rid, old_val, new_val in zip(rids, old_vals, new_vals):
+        if not values_equal(old_val, new_val):
+            store.append_diff(
+                step_id=step_id,
+                row_id=int(rid),
+                col=key,
+                old_val=old_val,
+                new_val=new_val,
+                change_type=ChangeType.MODIFIED,
+            )
+
+
+def _infer_series_operation(series: TrackedSeries) -> str:
+    """
+    Infer the operation that produced this Series.
+
+    Best effort - returns generic name if unknown.
+    """
+    if hasattr(series, "_tp_last_op") and series._tp_last_op is not None:
+        return series._tp_last_op
+    return "transform"
+
+
+# ============ STRING ACCESSOR WRAPPERS ============
+
+# Store original accessors module-level for restore
+_original_str_methods = {}
+
+# Use a WeakKeyDictionary to track series references without modifying accessor
+_str_accessor_series_map = weakref.WeakKeyDictionary()
+
+
+def wrap_str_accessor():
+    """
+    Wrap StringMethods to track operations.
+
+    We wrap the individual methods rather than __init__ since pandas
+    doesn't allow adding new attributes to accessor instances.
+    """
+    global _original_str_methods
+    from pandas.core.strings.accessor import StringMethods
+
+    # Wrap common string methods to preserve TrackedSeries
+    for method_name in [
+        "lower",
+        "upper",
+        "strip",
+        "lstrip",
+        "rstrip",
+        "replace",
+        "slice",
+        "split",
+        "contains",
+        "startswith",
+        "endswith",
+        "len",
+        "extract",
+        "findall",
+        "cat",
+        "get",
+        "pad",
+        "center",
+        "ljust",
+        "rjust",
+        "zfill",
+        "wrap",
+        "title",
+        "capitalize",
+        "swapcase",
+        "normalize",
+    ]:
+        if hasattr(StringMethods, method_name):
+            _wrap_str_method(StringMethods, method_name)
+
+
+def _wrap_str_method(cls, method_name):
+    """Wrap a single string method to preserve TrackedSeries."""
+    global _original_str_methods
+
+    if method_name in _original_str_methods:
+        return  # Already wrapped
+
+    original = getattr(cls, method_name)
+    _original_str_methods[method_name] = original
+
+    @wraps(original)
+    def wrapped(self, *args, **kwargs):
+        result = original(self, *args, **kwargs)
+        # Try to get the originating series from the accessor's internal _orig
+        series = getattr(self, "_orig", None)
+        if isinstance(result, pd.Series) and isinstance(series, TrackedSeries):
+            tracked = TrackedSeries(result)
+            # Copy weakref directly to avoid creating strong reference
+            tracked._tp_source_df_ref = getattr(series, "_tp_source_df_ref", None)
+            tracked._tp_source_col = getattr(series, "_tp_source_col", None)
+            tracked._tp_source_rids = getattr(series, "_tp_source_rids", None)
+            tracked._tp_last_op = f"str.{method_name}"
+            return tracked
+        return result
+
+    setattr(cls, method_name, wrapped)
+
+
+def instrument_series():
+    """Install all Series instrumentation."""
+    wrap_series_extraction()
+    wrap_series_assignment()
+    wrap_str_accessor()
+    # Note: DateTime accessor (.dt) wrapping is not implemented.
+    # Most datetime operations don't require cell-level tracking.
+
+
+def uninstrument_series():
+    """Restore original Series behavior."""
+    global _original_str_methods
+
+    if hasattr(pd.DataFrame, "_tp_original_getitem_series"):
+        pd.DataFrame.__getitem__ = pd.DataFrame._tp_original_getitem_series
+        delattr(pd.DataFrame, "_tp_original_getitem_series")
+    if hasattr(pd.DataFrame, "_tp_original_setitem_series"):
+        pd.DataFrame.__setitem__ = pd.DataFrame._tp_original_setitem_series
+        delattr(pd.DataFrame, "_tp_original_setitem_series")
+
+    # Restore str methods
+    if _original_str_methods:
+        try:
+            from pandas.core.strings.accessor import StringMethods
+
+            for method_name, original in _original_str_methods.items():
+                setattr(StringMethods, method_name, original)
+            _original_str_methods.clear()
+        except ImportError:
+            pass