diff --git a/tools/graph_query.py b/tools/graph_query.py new file mode 100644 index 0000000..30175d5 --- /dev/null +++ b/tools/graph_query.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python3 +""" +Filtered graph query for graphify-out/graph.json. + +Reduces noise by excluding test nodes and helper-only symbols. +Supports BFS query and shortest-path modes. + +Usage: + python scripts/graph_query.py "job lifecycle" + python scripts/graph_query.py --path "job()" "_destroy_clone()" + python scripts/graph_query.py "transport" --context calls --context implements + python scripts/graph_query.py "lifecycle" --include-tests + python scripts/graph_query.py "build" --depth 3 --budget 4000 +""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from collections import deque +from pathlib import Path + +DEFAULT_EXCLUDE = [ + r"^tests/", + r"/test_", + r"_test\.py$", + r"\.Tests\.", + r"Test-", + r"Measure-", +] + +SIGNAL_RELATIONS = { + "calls", + "implements", + "uses", + "defines", + "inherits", + "shares_data_with", + "references", +} + + +def load_graph(graph_path: Path) -> tuple[dict, dict]: + data = json.loads(graph_path.read_text(encoding="utf-8")) + nodes: dict[str, dict] = {n["id"]: n for n in data["nodes"]} + adj: dict[str, list[tuple[str, str, float]]] = {nid: [] for nid in nodes} + for link in data["links"]: + src, tgt = link["source"], link["target"] + rel = link.get("relation", "") + cs = link.get("confidence_score", 1.0) + if src in adj: + adj[src].append((tgt, rel, cs)) + if tgt in adj: + adj[tgt].append((src, rel, cs)) + return nodes, adj + + +def allowed_set(nodes: dict, exclude_patterns: list[str]) -> set[str]: + allowed: set[str] = set() + for nid, node in nodes.items(): + sf = node.get("source_file") or "" + if not any(re.search(pat, sf) for pat in exclude_patterns): + allowed.add(nid) + return allowed + + +def find_starts(nodes: dict, query: str, allowed: set[str]) -> list[str]: + q = query.lower() + exact, partial = [], [] + for nid in allowed: + label = (nodes[nid].get("label") or "").lower() + sf = (nodes[nid].get("source_file") or "").lower() + if q == label or q == nid: + exact.append(nid) + elif q in label or q in sf: + partial.append(nid) + return exact or partial + + +def bfs( + nodes: dict, + adj: dict, + starts: list[str], + allowed: set[str], + max_depth: int, + relations: set[str] | None, +) -> list[tuple[dict, int, str | None]]: + visited: dict[str, str | None] = {} + queue: deque[tuple[str, int, str | None]] = deque() + for nid in starts: + if nid in allowed: + visited[nid] = None + queue.append((nid, 0, None)) + + results: list[tuple[dict, int, str | None]] = [] + while queue: + nid, depth, via_rel = queue.popleft() + results.append((nodes[nid], depth, via_rel)) + if depth < max_depth: + for neighbor, rel, _cs in adj.get(nid, []): + if neighbor not in visited and neighbor in allowed: + if relations is None or rel in relations: + visited[neighbor] = nid + queue.append((neighbor, depth + 1, rel)) + return results + + +def shortest_path( + nodes: dict, + adj: dict, + a_ids: list[str], + b_ids: list[str], + allowed: set[str], +) -> list[tuple[str, str]]: + b_set = set(b_ids) + for start in a_ids: + if start not in allowed: + continue + visited: dict[str, tuple[str, str] | None] = {start: None} + queue: deque[str] = deque([start]) + while queue: + cur = queue.popleft() + if cur in b_set: + path: list[tuple[str, str]] = [] + node: str | None = cur + rel_used = "" + while node is not None: + prev = visited[node] + if prev is None: + path.append((node, "")) + else: + prev_node, rel_used = prev + path.append((node, rel_used)) + node = prev_node + continue + break + return list(reversed(path)) + for neighbor, rel, _cs in adj.get(cur, []): + if neighbor not in visited and neighbor in allowed: + visited[neighbor] = (cur, rel) + queue.append(neighbor) + return [] + + +def fmt_node(node: dict, depth: int, via_rel: str | None) -> str: + indent = " " * depth + label = node.get("label", node["id"]) + sf = node.get("source_file", "") + loc = node.get("source_location", "") + loc_str = f" loc={loc}" if loc else "" + rel_str = f" ←{via_rel}" if via_rel else "" + return f"{indent}NODE {label} [src={sf}{loc_str}]{rel_str}" + + +def main() -> None: + ap = argparse.ArgumentParser(description="Filtered graphify query") + ap.add_argument("query", nargs="?", help="Search term for BFS query") + ap.add_argument("--path", nargs=2, metavar=("A", "B"), help="Shortest path between two nodes") + ap.add_argument("--context", action="append", metavar="REL", help="Limit traversal to these edge relations (repeatable)") + ap.add_argument("--depth", type=int, default=2, help="BFS max depth (default 2)") + ap.add_argument("--budget", type=int, default=3000, help="Approximate output char budget (default 3000)") + ap.add_argument("--include-tests", action="store_true", help="Do not exclude test nodes") + ap.add_argument("--graph", type=Path, default=Path("graphify-out/graph.json"), help="Path to graph.json") + args = ap.parse_args() + + if not args.query and not args.path: + ap.print_help() + sys.exit(1) + + if not args.graph.exists(): + print(f"error: {args.graph} not found. Run /graphify first.", file=sys.stderr) + sys.exit(1) + + nodes, adj = load_graph(args.graph) + exclude = [] if args.include_tests else DEFAULT_EXCLUDE + allowed = allowed_set(nodes, exclude) + relations = set(args.context) if args.context else None + + total = len(nodes) + filtered_out = total - len(allowed) + print(f"Graph: {total} nodes, {filtered_out} excluded ({filtered_out*100//total}% noise filtered)\n") + + if args.path: + a_ids = find_starts(nodes, args.path[0], allowed) + b_ids = find_starts(nodes, args.path[1], allowed) + if not a_ids: + print(f"No match for '{args.path[0]}'", file=sys.stderr) + sys.exit(1) + if not b_ids: + print(f"No match for '{args.path[1]}'", file=sys.stderr) + sys.exit(1) + path = shortest_path(nodes, adj, a_ids, b_ids, allowed) + if not path: + print(f"No path found between '{args.path[0]}' and '{args.path[1]}'") + else: + print(f"Path ({len(path)} hops):") + for i, (nid, rel) in enumerate(path): + node = nodes[nid] + label = node.get("label", nid) + sf = node.get("source_file", "") + loc = node.get("source_location", "") + loc_str = f":{loc}" if loc else "" + arrow = f" --{rel}--> " if rel else "" + print(f" {' ' * i}{arrow}{label} [{sf}{loc_str}]") + return + + starts = find_starts(nodes, args.query, allowed) + if not starts: + print(f"No matching nodes for '{args.query}' (excluding tests)") + sys.exit(0) + + print(f"BFS from {len(starts)} start node(s) · depth={args.depth}" + + (f" · context={','.join(sorted(relations))}" if relations else "") + "\n") + + results = bfs(nodes, adj, starts, allowed, args.depth, relations) + + chars = 0 + shown = 0 + for node, depth, via_rel in results: + line = fmt_node(node, depth, via_rel) + if chars + len(line) > args.budget: + print(f"\n... ({len(results) - shown} more nodes, increase --budget to see)") + break + print(line) + chars += len(line) + 1 + shown += 1 + + +if __name__ == "__main__": + main()