coltrace 0.1.0.dev4__py3-none-any.whl

coltrace/history.py

@@ -0,0 +1,441 @@
+"""A column's recorded history, separate from its value dependency graph.
+
+Value edges answer whether one value can affect another. History also includes
+allocations, row filtering, callback statements, overwrites and destinations.
+Only explicit recorder facts establish execution order; source line numbers do
+not establish an ordering between functions or call contexts.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict, deque
+from dataclasses import asdict
+from typing import Any
+
+from .lineage_models import Column, LineageReport
+
+_SOURCE_KINDS = {"source", "allocation", "parameter", "external", "create_dataframe"}
+_USE_KINDS = {"read", "consume", "consumption", "sink", "write_file", "return"}
+_BOUNDARY_KINDS = {"opaque", "boundary", "unresolved", "opaque_call"}
+_COMPLETENESS = {"unresolved": 0, "partial": 1, "complete": 2}
+
+
+def lifetime_versions(report: LineageReport, selected: Column) -> set[str]:
+    """Versions in one label lifetime; an inplace rename ends that lifetime.
+
+    Recreating ``a`` after renaming it to ``renamed`` creates an independent
+    label. A shared frame ID and label alone cannot reconnect those values.
+    """
+    columns = {column.id: column for column in report.columns}
+    created = {
+        identifier: event.sequence for event in report.history for identifier in event.outputs
+    }
+    selected_sequence = created.get(selected.id)
+    vacated = [
+        event.sequence
+        for event in report.history
+        if event.kind == "rename"
+        and any(
+            identifier in columns
+            and (columns[identifier].dataframe_id, columns[identifier].name)
+            == (selected.dataframe_id, selected.name)
+            for identifier in event.inputs
+        )
+        and any(
+            identifier in columns
+            and columns[identifier].dataframe_id == selected.dataframe_id
+            and columns[identifier].name != selected.name
+            for identifier in event.outputs
+        )
+    ]
+    lower = max(
+        (
+            sequence
+            for sequence in vacated
+            if selected_sequence is not None and sequence < selected_sequence
+        ),
+        default=-1,
+    )
+    upper = min(
+        (
+            sequence
+            for sequence in vacated
+            if selected_sequence is not None and sequence >= selected_sequence
+        ),
+        default=float("inf"),
+    )
+    return {
+        column.id
+        for column in report.columns
+        if (column.dataframe_id, column.name) == (selected.dataframe_id, selected.name)
+        and (column.id not in created or lower < created[column.id] <= upper)
+    }
+
+
+def _compact_stages(stages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Collapse bookkeeping at one statement; retain raw event IDs for the DAG."""
+    result: list[dict[str, Any]] = []
+    groups: dict[tuple[str, str, str, int, str], dict[str, Any]] = {}
+    row_kinds = {"filter", "dropna", "projection"}
+    for stage in stages:
+        if stage["kind"] == "read" and any(
+            other["kind"] != "read"
+            and other["location"]["path"] == stage["location"]["path"]
+            and other["location"]["line"] == stage["location"]["line"]
+            and other["context"] == stage["context"]
+            and set(stage["inputs"]).intersection((*other["inputs"], *other["outputs"]))
+            for other in stages
+        ):
+            continue
+        if stage["kind"] == "consume" and stage["label"] in {
+            "argument",
+            "apply input",
+            "callback input",
+            "string transform",
+        }:
+            continue
+        kind = stage["kind"]
+        key_kind = "rows" if kind in row_kinds else kind
+        if kind == "bind" and any(
+            other["kind"] in row_kinds
+            and other["location"]["path"] == stage["location"]["path"]
+            and other["location"]["line"] == stage["location"]["line"]
+            and other["context"] == stage["context"]
+            for other in stages
+        ):
+            key_kind = "rows"
+        key = (
+            stage["context"],
+            stage["location"]["path"],
+            stage["scope"],
+            stage["location"]["line"],
+            key_kind,
+        )
+        # Separate callback assignments and positional return components even on
+        # a shared source line. Their branch identities matter.
+        merge = key_kind in {"rows", "sink", "boundary"}
+        previous = groups.get(key) if merge else None
+        if previous is None:
+            stage = {**stage, "event_ids": [stage["id"]]}
+            if key_kind == "rows":
+                stage["label"] = (
+                    "Project selected columns; preserve their values"
+                    if kind == "projection"
+                    else "Filter rows; preserve values of retained columns"
+                )
+            result.append(stage)
+            if merge:
+                groups[key] = stage
+        else:
+            previous["event_ids"].append(stage["id"])
+            for field in ("inputs", "outputs", "dataframe_ids", "parents"):
+                previous[field] = list(dict.fromkeys((*previous[field], *stage[field])))
+            if kind == "sink" and stage["label"] != "sink":
+                previous["label"], previous["expression"] = stage["label"], stage["expression"]
+            elif kind == "bind":
+                previous["expression"] = stage["expression"]
+    return result
+
+
+def _category(kind: str) -> str:
+    if kind in {"callback_control", "callback_configuration", "callback_context"}:
+        return kind.removeprefix("callback_")
+    if kind in _SOURCE_KINDS or kind.startswith("source"):
+        return "source"
+    if kind in _USE_KINDS or kind.startswith("sink"):
+        return "use"
+    if kind in _BOUNDARY_KINDS or "opaque" in kind or "unresolved" in kind:
+        return "boundary"
+    return "transform"
+
+
+def _fallback_events(report: LineageReport) -> list[dict[str, Any]]:
+    """Preserve usable old reports without claiming unrecorded inner steps."""
+    columns = {column.id: column for column in report.columns}
+    reads: dict[str, list[str]] = defaultdict(list)
+    writes: dict[str, list[str]] = defaultdict(list)
+    for edge in report.edges:
+        if edge.kind == "read":
+            reads[edge.target].append(edge.source)
+        elif edge.kind == "write":
+            writes[edge.source].append(edge.target)
+    events: list[dict[str, Any]] = []
+    for frame in report.dataframes:
+        events.append(
+            {
+                "id": f"history-source:{frame.id}",
+                "sequence": None,
+                "kind": "parameter" if "parameter" in frame.origin else "source",
+                "label": frame.origin,
+                "location": asdict(frame.location),
+                "module": frame.module,
+                "scope": frame.scope,
+                "context": frame.context,
+                "dataframe_ids": [frame.id],
+                "inputs": [],
+                "outputs": [],
+                "parents": [],
+                "expression": "",
+                "evidence": frame.evidence,
+                "completeness": "partial",
+            }
+        )
+    for operation in report.operations:
+        inputs = sorted(reads[operation.id])
+        outputs = sorted(writes[operation.id])
+        events.append(
+            {
+                **asdict(operation),
+                "sequence": None,
+                "label": operation.kind,
+                "dataframe_ids": sorted(
+                    {columns[item].dataframe_id for item in inputs + outputs if item in columns}
+                ),
+                "inputs": inputs,
+                "outputs": outputs,
+                "parents": [],
+                "expression": "",
+            }
+        )
+    for index, consumption in enumerate(report.consumptions):
+        events.append(
+            {
+                **asdict(consumption),
+                "id": f"history-use:{index}:{consumption.column_id}",
+                "sequence": None,
+                "kind": "consume",
+                "label": consumption.kind,
+                "dataframe_ids": [columns[consumption.column_id].dataframe_id],
+                "inputs": [consumption.column_id],
+                "outputs": [],
+                "parents": [],
+                "expression": "",
+                "evidence": "confirmed",
+                "completeness": "partial",
+            }
+        )
+    return events
+
+
+def _reachable_columns(report: LineageReport, selected_id: str) -> set[str]:
+    """Ancestors and descendants, without following unrelated sibling branches."""
+    reads: dict[str, set[str]] = defaultdict(set)
+    writes: dict[str, set[str]] = defaultdict(set)
+    for edge in report.edges:
+        if edge.kind == "read":
+            reads[edge.target].add(edge.source)
+        elif edge.kind == "write":
+            writes[edge.source].add(edge.target)
+    forward: dict[str, set[str]] = defaultdict(set)
+    reverse: dict[str, set[str]] = defaultdict(set)
+    for operation_id, inputs in reads.items():
+        for source in inputs:
+            for target in writes[operation_id]:
+                forward[source].add(target)
+                reverse[target].add(source)
+    result = {selected_id}
+    for adjacent in (forward, reverse):
+        queue = deque([selected_id])
+        seen = {selected_id}
+        while queue:
+            for target in adjacent[queue.popleft()]:
+                if target not in seen:
+                    seen.add(target)
+                    queue.append(target)
+        result.update(seen)
+    return result
+
+
+def build_history(
+    report: LineageReport,
+    selected_id: str,
+    included_ids: set[str] | None = None,
+    *,
+    evidence: str = "all",
+) -> dict[str, Any]:
+    """Return an event DAG and ordered narrative for a recorded column lifetime.
+
+    ``included_ids`` is the value-query slice. Earlier versions of the selected
+    logical column remain visible as history even when a constant overwrite
+    correctly has no dependency edge. A ``supersedes`` edge records chronology
+    only and must never be interpreted as a value dependency.
+    """
+    columns = {column.id: column for column in report.columns}
+    if selected_id not in columns:
+        raise ValueError(f"Unknown column ID: {selected_id}")
+    selected = columns[selected_id]
+    lifetime = lifetime_versions(report, selected)
+    included = (
+        _reachable_columns(report, selected_id) if included_ids is None else set(included_ids)
+    )
+    included.add(selected_id)
+    included.update(lifetime)
+    included.intersection_update(columns)
+    frame_ids = {columns[item].dataframe_id for item in included}
+    recorded = getattr(report, "history", ())
+    events = [asdict(event) for event in recorded] if recorded else _fallback_events(report)
+    if evidence == "confirmed":
+        events = [event for event in events if event["evidence"] == "confirmed"]
+    by_event = {event["id"]: event for event in events}
+    cutoffs: dict[str, int] = {}
+    event_ids: set[str] = set()
+    disconnected: set[str] = set()
+    for event in events:
+        touches = bool(
+            included.intersection(event["outputs"] if event["outputs"] else event["inputs"])
+        ) and not event["kind"].startswith("callback_")
+        related_frames = frame_ids.intersection(event["dataframe_ids"])
+        sequence = event.get("sequence")
+        beyond = sequence is not None and any(
+            frame in cutoffs and sequence > cutoffs[frame] for frame in related_frames
+        )
+        frame_only = (
+            (_category(event["kind"]) == "source" and not event["outputs"])
+            or event["kind"] in {"bind", "call", "return", "sink"}
+            or (event["kind"] == "boundary" and not event["inputs"] and not event["outputs"])
+        )
+        if touches or (related_frames and frame_only and not beyond):
+            event_ids.add(event["id"])
+            if _category(event["kind"]) == "boundary" and sequence is not None:
+                for frame in related_frames:
+                    cutoffs.setdefault(frame, sequence)
+        elif related_frames and event["kind"] == "sink" and beyond:
+            event_ids.add(event["id"])
+            disconnected.add(event["id"])
+    for event in events:
+        if event["id"] in disconnected:
+            event.update(
+                label="Later export of same DataFrame; column survival is unverified",
+                inputs=[],
+                outputs=[],
+                parents=[],
+                evidence="unresolved",
+                completeness="unresolved",
+                continuity="unresolved",
+            )
+    queue = deque(sorted(event_ids))
+    while queue:
+        event = by_event[queue.popleft()]
+        for parent in event.get("parents", ()):
+            if parent in by_event and parent not in event_ids:
+                event_ids.add(parent)
+                queue.append(parent)
+    chosen = [event for event in events if event["id"] in event_ids]
+    chosen.sort(
+        key=lambda event: (
+            event.get("sequence") is None,
+            event.get("sequence") if event.get("sequence") is not None else 0,
+            event["id"],
+        )
+    )
+    # Co-inputs explain an included transformation but do not pull in unrelated
+    # downstream branches of those co-inputs.
+    displayed_columns = included | {
+        column_id for event in chosen for column_id in event["inputs"] if column_id in columns
+    }
+    nodes: list[dict[str, Any]] = [
+        {
+            **asdict(columns[column_id]),
+            "node_type": "column",
+            "selected": column_id == selected_id,
+            "context_only": column_id not in included,
+        }
+        for column_id in sorted(displayed_columns)
+    ]
+    stages: list[dict[str, Any]] = []
+    edges: list[dict[str, Any]] = []
+    for event in chosen:
+        stage = {**event, "category": _category(event["kind"])}
+        stages.append(stage)
+        nodes.append({**stage, "node_type": "event", "selected": False})
+        for source in event["inputs"]:
+            if source in displayed_columns:
+                edges.append(
+                    {
+                        "source": source,
+                        "target": event["id"],
+                        "kind": "input",
+                        "evidence": event["evidence"],
+                    }
+                )
+        for target in event["outputs"]:
+            if target in displayed_columns:
+                edges.append(
+                    {
+                        "source": event["id"],
+                        "target": target,
+                        "kind": "output",
+                        "evidence": event["evidence"],
+                    }
+                )
+        for parent in event.get("parents", ()):
+            if parent in event_ids:
+                edges.append(
+                    {
+                        "source": parent,
+                        "target": event["id"],
+                        "kind": "step",
+                        "evidence": event["evidence"],
+                    }
+                )
+    versions = sorted(
+        (column for column in report.columns if column.id in lifetime and column.id in included),
+        key=lambda column: column.version,
+    )
+    for previous, current in zip(versions, versions[1:], strict=False):
+        creators = [event for event in chosen if current.id in event["outputs"]]
+        if creators and previous.id not in creators[0]["inputs"]:
+            edges.append(
+                {
+                    "source": previous.id,
+                    "target": creators[0]["id"],
+                    "kind": "supersedes",
+                    "evidence": creators[0]["evidence"],
+                }
+            )
+    source_events = [stage for stage in stages if stage["category"] == "source"]
+    # A frame allocation precedes observations of its external schema. This is
+    # origin metadata, not a claim that a read_csv expression defines a column.
+    for event in source_events:
+        for column_id in sorted(included):
+            column = columns[column_id]
+            if (
+                column.dataframe_id in event["dataframe_ids"]
+                and column.origin in {"external", "external_unknown", "parameter"}
+                and not any(edge["target"] == column_id for edge in edges)
+            ):
+                edges.append(
+                    {
+                        "source": event["id"],
+                        "target": column_id,
+                        "kind": "observed_column",
+                        "evidence": column.evidence,
+                    }
+                )
+    stages = _compact_stages(stages)
+    groups = {
+        f"{category}s": [stage for stage in stages if stage["category"] == category]
+        for category in ("source", "transform", "use", "boundary")
+    }
+    # Correct plural for the public narrative group.
+    groups["boundaries"] = groups.pop("boundarys")
+    completeness = min(
+        [columns[item].completeness for item in included]
+        + [event["completeness"] for event in chosen]
+        + (["partial"] if not recorded else ["complete"]),
+        key=_COMPLETENESS.__getitem__,
+    )
+    if completeness == "unresolved" and any(
+        columns[item].completeness != "unresolved" for item in included
+    ):
+        completeness = "partial"
+    return {
+        "selected_column_id": selected_id,
+        "nodes": nodes,
+        "edges": edges,
+        "stages": stages,
+        **groups,
+        "completeness": completeness,
+        "ordering": "recorded" if recorded else "unavailable",
+        "recorded": bool(recorded),
+    }