Add weighted affected traversal

graphify/affected.py

@@ -2,8 +2,9 @@
 
 from collections import deque
 from dataclasses import dataclass
+import heapq
 from pathlib import Path
-from typing import Iterable
+from typing import Iterable, Mapping
 import unicodedata
 
 import networkx as nx
@@ -23,6 +24,20 @@
     "embeds",
 )
 
+DEFAULT_RELATION_WEIGHTS: Mapping[str, float] = {
+    "calls": 1.0,
+    "references": 0.9,
+    "imports": 0.85,
+    "imports_from": 0.85,
+    "re_exports": 0.8,
+    "inherits": 0.75,
+    "extends": 0.75,
+    "implements": 0.75,
+    "uses": 1.0,
+    "mixes_in": 0.9,
+    "embeds": 0.7,
+}
+
 
 @dataclass(frozen=True)
 class AffectedHit:
@@ -31,6 +46,28 @@ class AffectedHit:
     via_relation: str
 
 
+@dataclass(frozen=True)
+class WeightedAffectedHit:
+    node_id: str
+    cost: float
+    via_relation: str
+    path: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class WeightedAffectedResult:
+    hits: tuple[WeightedAffectedHit, ...]
+    proof_paths: dict[str, tuple[str, ...]]
+    metrics: dict[str, int | float]
+
+
+@dataclass(frozen=True)
+class PreparedAffectedGraph:
+    graph: nx.Graph
+    incoming: dict[str, tuple[tuple[str, str, str], ...]]
+    degree: dict[str, int]
+
+
 def _node_label(graph: nx.Graph, node_id: str) -> str:
     data = graph.nodes[node_id]
     return str(data.get("label") or node_id)
@@ -93,6 +130,185 @@ def resolve_seed(graph: nx.Graph, query: str) -> str | None:
     return None
 
 
+def prepare_affected_graph(graph: nx.Graph) -> PreparedAffectedGraph:
+    incoming: dict[str, list[tuple[str, str, str]]] = {}
+    for source, target, data in graph.edges(data=True):
+        source_id = str(source)
+        target_id = str(target)
+        relation = str(data.get("relation", ""))
+        incoming.setdefault(target_id, []).append((source_id, target_id, relation))
+    ordered = {
+        node_id: tuple(sorted(edges, key=lambda edge: (edge[0], edge[2])))
+        for node_id, edges in incoming.items()
+    }
+    degree = {str(node_id): int(deg) for node_id, deg in graph.degree()}
+    return PreparedAffectedGraph(graph=graph, incoming=ordered, degree=degree)
+
+
+def _path_to_seed(parent: Mapping[str, tuple[str, str]], seed: str, node_id: str) -> tuple[str, ...]:
+    path = [node_id]
+    current = node_id
+    for _ in range(500):
+        if current == seed:
+            return tuple(path)
+        nxt = parent.get(current)
+        if not nxt:
+            return tuple()
+        current = nxt[0]
+        path.append(current)
+    return tuple()
+
+
+def weighted_affected_details(
+    graph: nx.Graph,
+    seed: str,
+    *,
+    relations: Iterable[str] = DEFAULT_AFFECTED_RELATIONS,
+    relation_weights: Mapping[str, float] | None = None,
+    max_cost: float | None = None,
+    max_nodes: int = 200,
+    hub_degree: int | None = None,
+    hub_penalty: float = 2.0,
+    expand_hubs: bool = False,
+    proof_targets: Iterable[str] = (),
+) -> WeightedAffectedResult:
+    relation_set = set(relations)
+    weights = dict(DEFAULT_RELATION_WEIGHTS)
+    if relation_weights:
+        weights.update({str(k): float(v) for k, v in relation_weights.items()})
+    prepared = prepare_affected_graph(graph)
+    limit = max(1, int(max_nodes))
+    max_allowed_cost = float("inf") if max_cost is None else max(0.0, float(max_cost))
+    hub_threshold = None if hub_degree is None else max(1, int(hub_degree))
+    penalty = max(0.0, float(hub_penalty))
+
+    dist: dict[str, float] = {seed: 0.0}
+    parent: dict[str, tuple[str, str]] = {}
+    via: dict[str, str] = {}
+    queue: list[tuple[float, int, str]] = [(0.0, 0, seed)]
+    counter = 1
+    visited: set[str] = set()
+    hits: list[WeightedAffectedHit] = []
+    traversed_edges = 0
+    hub_skips = 0
+    max_seen_cost = 0.0
+
+    while queue and len(hits) < limit:
+        cost, _order, current = heapq.heappop(queue)
+        if current in visited:
+            continue
+        if cost != dist.get(current):
+            continue
+        visited.add(current)
+        max_seen_cost = max(max_seen_cost, cost)
+        if current != seed:
+            path = _path_to_seed(parent, seed, current)
+            hits.append(
+                WeightedAffectedHit(
+                    node_id=current,
+                    cost=round(cost, 6),
+                    via_relation=via.get(current, ""),
+                    path=path,
+                )
+            )
+
+        is_hub = hub_threshold is not None and prepared.degree.get(current, 0) >= hub_threshold
+        if current != seed and is_hub and not expand_hubs:
+            hub_skips += 1
+            continue
+
+        for source, _target, relation in prepared.incoming.get(current, ()):
+            if relation not in relation_set:
+                continue
+            traversed_edges += 1
+            if source in visited:
+                continue
+            relation_cost = max(0.01, float(weights.get(relation, 1.0)))
+            source_is_hub = hub_threshold is not None and prepared.degree.get(source, 0) >= hub_threshold
+            next_cost = cost + relation_cost + (penalty if source_is_hub and source != seed else 0.0)
+            if next_cost > max_allowed_cost:
+                continue
+            if next_cost < dist.get(source, float("inf")):
+                dist[source] = next_cost
+                parent[source] = (current, relation)
+                via[source] = relation
+                heapq.heappush(queue, (next_cost, counter, source))
+                counter += 1
+
+    proof_paths = {
+        target: _path_to_seed(parent, seed, target)
+        for target in proof_targets
+        if target == seed or target in visited
+    }
+    return WeightedAffectedResult(
+        hits=tuple(hits),
+        proof_paths=proof_paths,
+        metrics={
+            "visited_nodes": len(visited),
+            "traversed_edges": traversed_edges,
+            "hub_skips": hub_skips,
+            "max_cost": round(max_seen_cost, 6),
+        },
+    )
+
+
+def weighted_affected_nodes(
+    graph: nx.Graph,
+    seed: str,
+    *,
+    relations: Iterable[str] = DEFAULT_AFFECTED_RELATIONS,
+    relation_weights: Mapping[str, float] | None = None,
+    max_cost: float | None = None,
+    max_nodes: int = 200,
+    hub_degree: int | None = None,
+    hub_penalty: float = 2.0,
+    expand_hubs: bool = False,
+) -> list[WeightedAffectedHit]:
+    return list(
+        weighted_affected_details(
+            graph,
+            seed,
+            relations=relations,
+            relation_weights=relation_weights,
+            max_cost=max_cost,
+            max_nodes=max_nodes,
+            hub_degree=hub_degree,
+            hub_penalty=hub_penalty,
+            expand_hubs=expand_hubs,
+        ).hits
+    )
+
+
+def affected_proof_path(
+    graph: nx.Graph,
+    seed: str,
+    target: str,
+    *,
+    relations: Iterable[str] = DEFAULT_AFFECTED_RELATIONS,
+    relation_weights: Mapping[str, float] | None = None,
+    max_cost: float | None = None,
+    max_nodes: int = 200,
+    hub_degree: int | None = None,
+    hub_penalty: float = 2.0,
+    expand_hubs: bool = False,
+) -> tuple[str, ...]:
+    if target == seed:
+        return (seed,)
+    details = weighted_affected_details(
+        graph,
+        seed,
+        relations=relations,
+        relation_weights=relation_weights,
+        max_cost=max_cost,
+        max_nodes=max_nodes,
+        hub_degree=hub_degree,
+        hub_penalty=hub_penalty,
+        expand_hubs=expand_hubs,
+        proof_targets=(target,),
+    )
+    return details.proof_paths.get(target, tuple())
+
+
 def affected_nodes(
     graph: nx.Graph,
     seed: str,

tests/test_weighted_affected.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+import networkx as nx
+
+from graphify.affected import (
+    affected_proof_path,
+    weighted_affected_details,
+    weighted_affected_nodes,
+)
+
+
+def test_weighted_affected_orders_by_relation_cost() -> None:
+    graph = nx.DiGraph()
+    graph.add_edge("fast_consumer", "core", relation="calls")
+    graph.add_edge("slow_importer", "core", relation="imports")
+    graph.add_edge("downstream", "slow_importer", relation="calls")
+
+    hits = weighted_affected_nodes(
+        graph,
+        "core",
+        relations=("calls", "imports"),
+        relation_weights={"calls": 0.2, "imports": 2.0},
+    )
+
+    assert [h.node_id for h in hits[:2]] == ["fast_consumer", "slow_importer"]
+    assert hits[0].cost < hits[1].cost
+    assert hits[0].path == ("fast_consumer", "core")
+
+
+def test_weighted_affected_applies_hub_penalty_and_pruning() -> None:
+    graph = nx.DiGraph()
+    graph.add_edge("hub", "core", relation="calls")
+    for i in range(5):
+        graph.add_edge(f"spoke_{i}", "hub", relation="calls")
+
+    details = weighted_affected_details(
+        graph,
+        "core",
+        hub_degree=3,
+        hub_penalty=5.0,
+    )
+
+    assert [h.node_id for h in details.hits] == ["hub"]
+    assert details.hits[0].cost == 6.0
+    assert details.metrics["hub_skips"] == 1
+
+
+def test_weighted_affected_exports_proof_path_to_target() -> None:
+    graph = nx.DiGraph()
+    graph.add_edge("controller", "service", relation="calls")
+    graph.add_edge("service", "core", relation="calls")
+    graph.add_edge("unrelated", "other", relation="calls")
+
+    path = affected_proof_path(graph, "core", "controller")
+
+    assert path == ("controller", "service", "core")