coltrace 0.1.0.dev4__py3-none-any.whl

coltrace/graph.py

@@ -0,0 +1,380 @@
+"""Column queries over the rich operation graph and its column projection."""
+
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import asdict
+from typing import Any
+
+import networkx as nx
+
+from .lineage_models import Column, LineageReport
+
+_CONFIDENCE = {"unresolved": 0, "inferred": 1, "confirmed": 2}
+_COMPLETENESS = {"unresolved": 0, "partial": 1, "complete": 2}
+
+
+class ColumnNotFoundError(ValueError):
+    """No column matched the requested identity and filters."""
+
+
+class AmbiguousColumnError(ValueError):
+    """Several DataFrame identities matched; callers must disambiguate."""
+
+    def __init__(self, selector: str, candidates: tuple[Column, ...]) -> None:
+        self.candidates = candidates
+        super().__init__(f"Column {selector!r} is ambiguous ({len(candidates)} identities).")
+
+
+def _weakest(values: list[str], ranks: dict[str, int]) -> str:
+    return min(values, key=ranks.__getitem__)
+
+
+class LineageGraph:
+    """Build queryable graphs without equating columns by their textual names.
+
+    ``rich`` contains column and operation nodes. ``columns`` projects each
+    read/write pair into one relation, preserving parallel operations. Selection
+    by name chooses the latest version *within each frame*; distinct frame IDs
+    remain ambiguous. Selecting an ID always refers to that exact version.
+    """
+
+    def __init__(self, report: LineageReport) -> None:
+        self.report = report
+        self.by_id = {column.id: column for column in report.columns}
+        self.rich: nx.MultiDiGraph[str] = nx.MultiDiGraph()
+        self.columns: nx.MultiDiGraph[str] = nx.MultiDiGraph()
+        for column in report.columns:
+            self.rich.add_node(column.id, node_type="column", **asdict(column))
+            self.columns.add_node(column.id, **asdict(column))
+        for operation in report.operations:
+            self.rich.add_node(operation.id, node_type="operation", **asdict(operation))
+        for edge in report.edges:
+            self.rich.add_edge(edge.source, edge.target, **asdict(edge))
+        for operation in report.operations:
+            reads = [
+                (source, data)
+                for source, _, data in self.rich.in_edges(operation.id, data=True)
+                if data["kind"] == "read" and source in self.by_id
+            ]
+            writes = [
+                (target, data)
+                for _, target, data in self.rich.out_edges(operation.id, data=True)
+                if data["kind"] == "write" and target in self.by_id
+            ]
+            relation = operation.kind if operation.kind in {"rename", "pass"} else "derive"
+            for source, read in reads:
+                for target, write in writes:
+                    evidence = _weakest(
+                        [
+                            operation.evidence,
+                            read["evidence"],
+                            write["evidence"],
+                            self.by_id[source].evidence,
+                            self.by_id[target].evidence,
+                        ],
+                        _CONFIDENCE,
+                    )
+                    self.columns.add_edge(
+                        source,
+                        target,
+                        kind=relation,
+                        operation_id=operation.id,
+                        location=asdict(operation.location),
+                        evidence=evidence,
+                        completeness=_weakest(
+                            [
+                                operation.completeness,
+                                self.by_id[source].completeness,
+                                self.by_id[target].completeness,
+                            ],
+                            _COMPLETENESS,
+                        ),
+                    )
+        self._passes: nx.DiGraph[str] = nx.DiGraph()
+        self._passes.add_nodes_from(self.by_id)
+        for source, target, data in self.columns.edges(data=True):
+            if data["kind"] == "pass" and self.by_id[source].name == self.by_id[target].name:
+                self._passes.add_edge(source, target)
+
+    def select(
+        self,
+        selector: str,
+        *,
+        dataframe: str | None = None,
+        module: str | None = None,
+        function: str | None = None,
+        version: int | None = None,
+    ) -> Column:
+        """Select exactly one column; filters use exact, case-sensitive names."""
+        candidates = [
+            column
+            for column in self.report.columns
+            if (column.id == selector if selector in self.by_id else column.name == selector)
+            and (dataframe is None or column.dataframe_id == dataframe)
+            and (module is None or column.module == module)
+            and (function is None or column.scope == function)
+            and (version is None or column.version == version)
+        ]
+        if selector not in self.by_id and version is None:
+            latest: dict[tuple[str, str], Column] = {}
+            for column in candidates:
+                key = (column.dataframe_id, column.name)
+                if key not in latest or column.version > latest[key].version:
+                    latest[key] = column
+            candidates = list(latest.values())
+        candidates.sort(key=lambda column: column.id)
+        if not candidates:
+            raise ColumnNotFoundError(f"No column matches {selector!r} and these filters.")
+        if len(candidates) > 1:
+            family = nx.node_connected_component(self._passes.to_undirected(), candidates[0].id)
+            if (
+                selector not in self.by_id
+                and version is None
+                and all(column.id in family for column in candidates)
+            ):
+                # Copies/projections of the same values are one selectable
+                # lineage family. A same-name value derivation or an unrelated
+                # DataFrame still requires explicit disambiguation.
+                candidates.sort(
+                    key=lambda column: (len(nx.ancestors(self._passes, column.id)), column.id)
+                )
+            else:
+                raise AmbiguousColumnError(selector, tuple(candidates))
+        return candidates[0]
+
+    def _filtered(self, relation: str | None, evidence: str) -> nx.MultiDiGraph[str]:
+        if relation not in {None, "derive", "rename", "pass"}:
+            raise ValueError("relation must be derive, rename or pass")
+        if evidence not in {"all", "confirmed"}:
+            raise ValueError("evidence must be all or confirmed")
+        graph: nx.MultiDiGraph[str] = nx.MultiDiGraph()
+        graph.add_nodes_from(self.columns.nodes(data=True))
+        for source, target, key, data in self.columns.edges(keys=True, data=True):
+            if (relation is None or data["kind"] == relation) and (
+                evidence == "all" or data["evidence"] == "confirmed"
+            ):
+                graph.add_edge(source, target, key=key, **data)
+        return graph
+
+    def _walk(self, graph: nx.MultiDiGraph[str], start: str, depth: int | None) -> dict[str, Any]:
+        # Track evidence states separately: a longer confirmed path must not be
+        # hidden by an earlier, shorter inferred path. Finite states stop cycles.
+        queue = deque([(start, "confirmed", "complete", 0)])
+        seen = {(start, "confirmed", "complete"): 0}
+        reached: dict[str, tuple[int, str, str]] = {}
+        while queue:
+            node, confidence, completeness, distance = queue.popleft()
+            if depth is not None and distance >= depth:
+                continue
+            for _, target, data in graph.out_edges(node, data=True):
+                new_confidence = _weakest([confidence, data["evidence"]], _CONFIDENCE)
+                new_completeness = _weakest([completeness, data["completeness"]], _COMPLETENESS)
+                next_distance = distance + 1
+                state = (target, new_confidence, new_completeness)
+                if state in seen and seen[state] <= next_distance:
+                    continue
+                seen[state] = next_distance
+                queue.append((target, new_confidence, new_completeness, next_distance))
+                if target == start:
+                    continue
+                previous = reached.get(target)
+                if previous is None:
+                    reached[target] = (next_distance, new_confidence, new_completeness)
+                else:
+                    best = max(
+                        [(previous[1], previous[2]), (new_confidence, new_completeness)],
+                        key=lambda item: (_CONFIDENCE[item[0]], _COMPLETENESS[item[1]]),
+                    )
+                    reached[target] = (min(previous[0], next_distance), *best)
+        direct: list[dict[str, Any]] = []
+        transitive: list[dict[str, Any]] = []
+        for node, (distance, confidence, completeness) in sorted(
+            reached.items(), key=lambda item: (item[1][0], item[0])
+        ):
+            entry = {
+                **asdict(self.by_id[node]),
+                "distance": distance,
+                "path_evidence": confidence,
+                "path_completeness": completeness,
+            }
+            (direct if distance == 1 else transitive).append(entry)
+        return {"direct": direct, "transitive": transitive}
+
+    def trace(
+        self,
+        column_id: str,
+        *,
+        direction: str = "both",
+        depth: int | None = None,
+        relation: str | None = None,
+        evidence: str = "all",
+    ) -> dict[str, Any]:
+        """Return provenance and potential consumers, counted in column hops.
+
+        Path evidence describes the strongest available path inside the chosen
+        depth, whose weakest relation determines that path's evidence. Distance
+        is independently the shortest eligible path. Impact is potential: an
+        edge does not establish that changed inputs change a consumer's output.
+        """
+        if column_id not in self.by_id:
+            raise ColumnNotFoundError(f"Unknown column ID: {column_id}")
+        if direction not in {"both", "upstream", "downstream"}:
+            raise ValueError("direction must be both, upstream or downstream")
+        if depth is not None and (type(depth) is not int or depth < 0):
+            raise ValueError("depth must be a non-negative integer")
+        graph = self._filtered(relation, evidence)
+        upstream = (
+            self._walk(graph.reverse(copy=False), column_id, depth)
+            if direction in {"both", "upstream"}
+            else {"direct": [], "transitive": []}
+        )
+        downstream = (
+            self._walk(graph, column_id, depth)
+            if direction in {"both", "downstream"}
+            else {"direct": [], "transitive": []}
+        )
+        included = {column_id}
+        for result in (upstream, downstream):
+            included.update(item["id"] for group in result.values() for item in group)
+        relations = [
+            {"source": source, "target": target, **data}
+            for source, target, data in graph.edges(data=True)
+            if source in included and target in included
+        ]
+        relations.sort(key=lambda item: (item["source"], item["target"], item["operation_id"]))
+        consumers = [item for item in self.report.consumptions if item.column_id in included]
+        relevant_operations = {item["operation_id"] for item in relations}
+        functions = {(self.by_id[node].module, self.by_id[node].scope) for node in included}
+        functions.update((item.module, item.scope) for item in consumers)
+        for operation in self.report.operations:
+            if operation.id in relevant_operations:
+                functions.add((operation.module, operation.scope))
+        from .history import build_history, lifetime_versions
+        from .journey import build_journey
+
+        # A latest-version value query cannot reach a derivative created before
+        # a constant overwrite. The lifetime view separately includes earlier
+        # versions, under the same graph filters. An explicitly older version
+        # does not acquire derivatives of a later independent reset.
+        selected = self.by_id[column_id]
+        journey_included = set(included)
+        versions = lifetime_versions(self.report, selected)
+        journey_seeds = {column_id}
+        for earlier in self.report.columns:
+            if (
+                earlier.id not in versions
+                or earlier.version > selected.version
+                or (evidence == "confirmed" and earlier.evidence != "confirmed")
+            ):
+                continue
+            journey_seeds.add(earlier.id)
+            journey_included.add(earlier.id)
+            for adjacent, eligible in (
+                (graph, direction in {"both", "downstream"}),
+                (graph.reverse(copy=False), direction in {"both", "upstream"}),
+            ):
+                if eligible:
+                    reached = self._walk(adjacent, earlier.id, depth)
+                    journey_included.update(
+                        item["id"] for group in reached.values() for item in group
+                    )
+        journey_relations = [
+            {"source": source, "target": target, **data}
+            for source, target, data in graph.edges(data=True)
+            if source in journey_included and target in journey_included
+        ]
+        history = build_history(self.report, column_id, journey_included, evidence=evidence)
+        journey = build_journey(
+            self.report,
+            column_id,
+            journey_included,
+            evidence=evidence,
+            relations=journey_relations,
+        )
+        journey["seed_column_ids"] = sorted(journey_seeds)
+        locations = []
+        for node in history["nodes"]:
+            place = node["location"]
+            if node.get("kind") == "call" or node.get("label") == "callback input":
+                place = {**place, "end_line": place["line"]}
+            locations.append((place, node["scope"]))
+        for item in consumers:
+            place = asdict(item.location)
+            if item.kind == "callback_input":
+                place["end_line"] = place["line"]
+            locations.append((place, item.scope))
+
+        def relevant_diagnostic(location: dict[str, Any], scope: str) -> bool:
+            # A diagnostic elsewhere in a large function is not a statement
+            # about this column. Recorded boundaries and exact source spans
+            # identify the relevant uncertainty instead of scope-wide noise.
+            return any(
+                scope == other_scope
+                and location["path"] == other["path"]
+                and location["line"] <= (other["end_line"] or other["line"])
+                and other["line"] <= (location["end_line"] or location["line"])
+                for other, other_scope in locations
+            )
+
+        return {
+            "column": asdict(self.by_id[column_id]),
+            "upstream": upstream,
+            "downstream": downstream,
+            "consumptions": [asdict(item) for item in consumers],
+            "functions": [
+                {"module": module, "scope": scope} for module, scope in sorted(functions)
+            ],
+            "diagnostics": [
+                asdict(item)
+                for item in self.report.diagnostics
+                if relevant_diagnostic(asdict(item.location), item.scope)
+                and not (
+                    item.code == "CT_UNSUPPORTED_CONTROL_FLOW"
+                    and any(
+                        str(node.get("kind", "")).startswith("callback_")
+                        and node.get("label", "").startswith("Type guard;")
+                        and node["location"]["path"] == item.location.path
+                        and node["location"]["line"] == item.location.line
+                        and node["scope"] == item.scope
+                        for node in history["nodes"]
+                    )
+                )
+            ],
+            "relations": relations,
+            "depth": depth,
+            "direction": direction,
+            "evidence": evidence,
+            "relation": relation,
+            "history": history,
+            "journey": journey,
+        }
+
+    def impact(self, column_id: str, **kwargs: Any) -> dict[str, Any]:
+        """List potential direct and transitive downstream impact."""
+        result = self.trace(column_id, direction="downstream", **kwargs)
+        direct = {(item["module"], item["scope"]) for item in result["downstream"]["direct"]}
+        transitive = {
+            (item["module"], item["scope"]) for item in result["downstream"]["transitive"]
+        }
+        direct.update(
+            (item["module"], item["scope"])
+            for item in result["consumptions"]
+            if item["column_id"] == column_id
+        )
+        transitive.update(
+            (item["module"], item["scope"])
+            for item in result["consumptions"]
+            if item["column_id"] != column_id
+        )
+        transitive.difference_update(direct)
+        result["direct_functions"] = [
+            {"module": module, "scope": scope} for module, scope in sorted(direct)
+        ]
+        result["transitive_functions"] = [
+            {"module": module, "scope": scope} for module, scope in sorted(transitive)
+        ]
+        result["functions"] = [
+            {"module": module, "scope": scope} for module, scope in sorted(direct | transitive)
+        ]
+        return result