_callgraph.py
python
sha256:133f9bcf57a62ec7ebf0cd71138b54a200b15989aea2fb2a6912497e2926aa8a
feat(dev-safety): Phase 1 of #185 — audit tool for ambiguou…
Sonnet 5
patch
2 days ago
| 1 | """Call-graph construction for the code-domain CLI commands. |
| 2 | |
| 3 | Provides two data structures built from the symbol graph and AST walk: |
| 4 | |
| 5 | ``ForwardGraph`` |
| 6 | ``caller_address → frozenset[callee_bare_name]`` |
| 7 | "What does this function call?" |
| 8 | |
| 9 | ``ReverseGraph`` |
| 10 | ``callee_bare_name → list[caller_address]`` |
| 11 | "What calls this function name?" |
| 12 | |
| 13 | Both structures cover **Python only** (stdlib ``ast``). tree-sitter languages |
| 14 | receive import-level analysis through the existing import symbol extraction; |
| 15 | call-site extraction requires a tree-sitter query extension that is deferred. |
| 16 | |
| 17 | The reverse graph is the foundation for three new commands: |
| 18 | |
| 19 | * ``muse impact`` — transitive blast-radius of a change |
| 20 | * ``muse dead`` — symbols with no callers and no importers |
| 21 | * ``muse coverage`` — method call-coverage for a class interface |
| 22 | |
| 23 | Design note |
| 24 | ----------- |
| 25 | All graph-building functions perform a single linear pass over the manifest |
| 26 | (read blob → parse AST → walk). They are intentionally not cached, so that |
| 27 | callers always see the committed state for a given manifest. CLI commands |
| 28 | that need the graph multiple times should call once and pass the result. |
| 29 | """ |
| 30 | |
| 31 | import ast |
| 32 | import logging |
| 33 | import pathlib |
| 34 | from collections import deque |
| 35 | |
| 36 | from muse.core.object_store import read_object |
| 37 | from muse.core.symbol_cache import SymbolCache |
| 38 | from muse.core.callgraph_cache import CallGraphCache |
| 39 | from muse.plugins.code.ast_parser import SymbolTree, parse_symbols |
| 40 | from muse.core.validation import MAX_AST_BYTES |
| 41 | from muse.core.types import Manifest |
| 42 | from muse.plugins.code._framework import ( # re-exported for caller convenience |
| 43 | ImplicitEntryEdge, |
| 44 | ImplicitEdgeGraph, |
| 45 | build_implicit_edge_graph, |
| 46 | ) |
| 47 | |
| 48 | logger = logging.getLogger(__name__) |
| 49 | |
| 50 | #: Mapping from caller symbol address to the set of bare callee names. |
| 51 | ForwardGraph = dict[str, frozenset[str]] |
| 52 | |
| 53 | #: Mapping from bare callee name to the list of caller addresses. |
| 54 | ReverseGraph = dict[str, list[str]] |
| 55 | |
| 56 | _PY_SUFFIXES: frozenset[str] = frozenset({".py", ".pyi"}) |
| 57 | |
| 58 | # --------------------------------------------------------------------------- |
| 59 | # AST helpers |
| 60 | # --------------------------------------------------------------------------- |
| 61 | |
| 62 | def call_name(func_node: ast.expr) -> str | None: |
| 63 | """Return the bare callee name from an ``ast.Call`` func node, or None. |
| 64 | |
| 65 | Handles simple names (``foo()``) and attribute access (``obj.method()``). |
| 66 | Ignores subscript calls and other exotic forms. |
| 67 | """ |
| 68 | if isinstance(func_node, ast.Name): |
| 69 | return func_node.id |
| 70 | if isinstance(func_node, ast.Attribute): |
| 71 | return func_node.attr |
| 72 | return None |
| 73 | |
| 74 | def find_func_node( |
| 75 | stmts: list[ast.stmt], |
| 76 | name_parts: list[str], |
| 77 | ) -> ast.FunctionDef | ast.AsyncFunctionDef | None: |
| 78 | """Recursively locate a function node by its dotted qualified-name parts. |
| 79 | |
| 80 | Args: |
| 81 | stmts: Statement list from a module or class body. |
| 82 | name_parts: Dotted path components, e.g. ``["User", "save"]``. |
| 83 | |
| 84 | Returns: |
| 85 | The matching ``FunctionDef``/``AsyncFunctionDef`` node, or ``None``. |
| 86 | """ |
| 87 | if not name_parts: |
| 88 | return None |
| 89 | target = name_parts[0] |
| 90 | for stmt in stmts: |
| 91 | if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)) and stmt.name == target: |
| 92 | if len(name_parts) == 1: |
| 93 | return stmt |
| 94 | elif isinstance(stmt, ast.ClassDef) and stmt.name == target and len(name_parts) > 1: |
| 95 | return find_func_node(stmt.body, name_parts[1:]) |
| 96 | return None |
| 97 | |
| 98 | def callees_for_symbol(source: bytes, address: str) -> list[str]: |
| 99 | """Return sorted unique bare callee names called by the Python symbol at *address*. |
| 100 | |
| 101 | Args: |
| 102 | source: Raw bytes of the Python source file. |
| 103 | address: Full symbol address, e.g. ``"src/billing.py::compute_invoice_total"``. |
| 104 | |
| 105 | Returns: |
| 106 | Sorted list of bare callee names. Empty if the file has a syntax |
| 107 | error, the symbol is not found, or it is not a function/method. |
| 108 | """ |
| 109 | if "::" not in address: |
| 110 | return [] |
| 111 | sym_qualified = address.split("::", 1)[1] |
| 112 | if len(source) > MAX_AST_BYTES: |
| 113 | return [] |
| 114 | try: |
| 115 | tree = ast.parse(source) |
| 116 | except SyntaxError: |
| 117 | return [] |
| 118 | func_node = find_func_node(tree.body, sym_qualified.split(".")) |
| 119 | if func_node is None: |
| 120 | return [] |
| 121 | names: set[str] = set() |
| 122 | for node in ast.walk(func_node): |
| 123 | if isinstance(node, ast.Call): |
| 124 | name = call_name(node.func) |
| 125 | if name: |
| 126 | names.add(name) |
| 127 | return sorted(names) |
| 128 | |
| 129 | # --------------------------------------------------------------------------- |
| 130 | # Graph construction |
| 131 | # --------------------------------------------------------------------------- |
| 132 | |
| 133 | def build_forward_graph( |
| 134 | root: pathlib.Path, |
| 135 | manifest: Manifest, |
| 136 | cache: SymbolCache | None = None, |
| 137 | callgraph_cache: CallGraphCache | None = None, |
| 138 | ) -> ForwardGraph: |
| 139 | """Build the forward call graph from *manifest*. |
| 140 | |
| 141 | Scans every Python file in the manifest, walking each symbol's body for |
| 142 | ``ast.Call`` nodes. Two levels of caching are supported: |
| 143 | |
| 144 | * ``callgraph_cache`` — fastest path: skips ``read_object``, ``ast.parse``, |
| 145 | and the entire AST walk for files whose per-file subgraph is already |
| 146 | cached. On a warm cache, each file costs only a dict lookup and a |
| 147 | ``graph.update()`` call. |
| 148 | * ``cache`` (``SymbolCache``) — skips ``parse_symbols`` on a hit, but |
| 149 | ``read_object`` and ``ast.parse`` are still required. |
| 150 | |
| 151 | The caller is responsible for calling ``callgraph_cache.save()`` after |
| 152 | the build if persistence is desired. |
| 153 | |
| 154 | Args: |
| 155 | root: Repository root (used to read blobs from the object store). |
| 156 | manifest: Snapshot manifest mapping file path → SHA-256 object ID. |
| 157 | cache: Optional shared ``SymbolCache`` instance. |
| 158 | callgraph_cache: Optional ``CallGraphCache`` instance. When provided |
| 159 | and a file's subgraph is cached, the file is skipped |
| 160 | entirely — no blob fetch, no AST parse. |
| 161 | |
| 162 | Returns: |
| 163 | ``{caller_address: frozenset[callee_bare_name]}``. |
| 164 | """ |
| 165 | graph: ForwardGraph = {} |
| 166 | for file_path, obj_id in manifest.items(): |
| 167 | if pathlib.PurePosixPath(file_path).suffix.lower() not in _PY_SUFFIXES: |
| 168 | continue |
| 169 | |
| 170 | # Fast path: return cached per-file subgraph — no read_object, no ast.parse. |
| 171 | if callgraph_cache is not None: |
| 172 | cached = callgraph_cache.get(obj_id) |
| 173 | if cached is not None: |
| 174 | graph.update(cached) |
| 175 | continue |
| 176 | |
| 177 | raw = read_object(root, obj_id) |
| 178 | if raw is None: |
| 179 | continue |
| 180 | if len(raw) > MAX_AST_BYTES: |
| 181 | continue |
| 182 | try: |
| 183 | ast_tree = ast.parse(raw) |
| 184 | except SyntaxError: |
| 185 | continue |
| 186 | # Use the symbol cache for symbol tree if available; fall back to direct parse. |
| 187 | sym_tree: SymbolTree = ( |
| 188 | cache.get(obj_id) or parse_symbols(raw, file_path) |
| 189 | if cache is not None |
| 190 | else parse_symbols(raw, file_path) |
| 191 | ) |
| 192 | file_subgraph: ForwardGraph = {} |
| 193 | for addr, rec in sym_tree.items(): |
| 194 | if rec["kind"] not in {"function", "async_function", "method", "async_method"}: |
| 195 | continue |
| 196 | func_node = find_func_node(ast_tree.body, rec["qualified_name"].split(".")) |
| 197 | if func_node is None: |
| 198 | continue |
| 199 | names: set[str] = set() |
| 200 | for node in ast.walk(func_node): |
| 201 | if isinstance(node, ast.Call): |
| 202 | name = call_name(node.func) |
| 203 | if name: |
| 204 | names.add(name) |
| 205 | file_subgraph[addr] = frozenset(names) |
| 206 | |
| 207 | if callgraph_cache is not None: |
| 208 | callgraph_cache.put(obj_id, file_subgraph) |
| 209 | |
| 210 | graph.update(file_subgraph) |
| 211 | return graph |
| 212 | |
| 213 | def build_reverse_graph( |
| 214 | root: pathlib.Path, |
| 215 | manifest: Manifest, |
| 216 | cache: SymbolCache | None = None, |
| 217 | callgraph_cache: CallGraphCache | None = None, |
| 218 | ) -> ReverseGraph: |
| 219 | """Build the reverse call graph from *manifest*. |
| 220 | |
| 221 | Inverts the forward graph: maps each callee bare name to every caller |
| 222 | address that calls it. |
| 223 | |
| 224 | Args: |
| 225 | root: Repository root. |
| 226 | manifest: Snapshot manifest mapping file path → SHA-256 object ID. |
| 227 | cache: Optional shared ``SymbolCache`` instance. |
| 228 | callgraph_cache: Optional ``CallGraphCache`` instance. When provided |
| 229 | and populated, skips ``read_object`` + ``ast.parse`` |
| 230 | for cached files. |
| 231 | |
| 232 | Returns: |
| 233 | ``{callee_bare_name: [caller_address, ...]}``. |
| 234 | """ |
| 235 | forward = build_forward_graph(root, manifest, cache=cache, callgraph_cache=callgraph_cache) |
| 236 | reverse: ReverseGraph = {} |
| 237 | for caller_addr, callee_names in forward.items(): |
| 238 | for name in callee_names: |
| 239 | reverse.setdefault(name, []).append(caller_addr) |
| 240 | # Sort each caller list for deterministic output. |
| 241 | for name in reverse: |
| 242 | reverse[name].sort() |
| 243 | return reverse |
| 244 | |
| 245 | def transitive_callers( |
| 246 | start_name: str, |
| 247 | reverse: ReverseGraph, |
| 248 | max_depth: int = 0, |
| 249 | ) -> dict[int, list[str]]: |
| 250 | """BFS through *reverse* to find all transitive callers of *start_name*. |
| 251 | |
| 252 | Args: |
| 253 | start_name: Bare function/method name to start from. |
| 254 | reverse: Reverse call graph produced by :func:`build_reverse_graph`. |
| 255 | max_depth: Maximum BFS depth. ``0`` means unlimited. |
| 256 | |
| 257 | Returns: |
| 258 | ``{depth: [caller_address, ...]}``, depth 1 = direct callers. |
| 259 | Addresses that appear at multiple depths are recorded only at the |
| 260 | shallowest depth (first encounter wins). |
| 261 | """ |
| 262 | result: dict[int, list[str]] = {} |
| 263 | visited_names: set[str] = {start_name} |
| 264 | visited_addrs: set[str] = set() |
| 265 | # Use deque for O(1) popleft instead of O(n) list.pop(0). |
| 266 | q: deque[tuple[str, int]] = deque([(start_name, 0)]) |
| 267 | |
| 268 | while q: |
| 269 | name, depth = q.popleft() |
| 270 | if max_depth > 0 and depth >= max_depth: |
| 271 | continue |
| 272 | next_depth = depth + 1 |
| 273 | for caller_addr in reverse.get(name, []): |
| 274 | if caller_addr in visited_addrs: |
| 275 | continue |
| 276 | visited_addrs.add(caller_addr) |
| 277 | result.setdefault(next_depth, []).append(caller_addr) |
| 278 | # Extract the bare name of the caller to continue BFS. |
| 279 | caller_name = caller_addr.split("::")[-1].split(".")[-1] |
| 280 | if caller_name not in visited_names: |
| 281 | visited_names.add(caller_name) |
| 282 | q.append((caller_name, next_depth)) |
| 283 | |
| 284 | return result |
| 285 | |
| 286 | def transitive_callees( |
| 287 | start_addr: str, |
| 288 | forward: ForwardGraph, |
| 289 | max_depth: int = 0, |
| 290 | ) -> dict[int, list[str]]: |
| 291 | """BFS through *forward* to find all transitive callees of *start_addr*. |
| 292 | |
| 293 | Args: |
| 294 | start_addr: Full caller address (``file::symbol``) or bare name to |
| 295 | start from. |
| 296 | forward: Forward call graph produced by :func:`build_forward_graph`. |
| 297 | max_depth: Maximum BFS depth. ``0`` means unlimited. |
| 298 | |
| 299 | Returns: |
| 300 | ``{depth: [callee_bare_name, ...]}``, depth 1 = direct callees. |
| 301 | Names that appear at multiple depths are recorded only at the |
| 302 | shallowest depth (first encounter wins). |
| 303 | """ |
| 304 | # Seed from the start address's own callee set if it exists; otherwise |
| 305 | # treat start_addr as a bare name and seed from any forward entry whose |
| 306 | # address ends with that name. |
| 307 | initial: frozenset[str] |
| 308 | if start_addr in forward: |
| 309 | initial = forward[start_addr] |
| 310 | else: |
| 311 | # Bare-name fallback: collect callees from any matching address. |
| 312 | bare = start_addr.split("::")[-1].split(".")[-1] |
| 313 | combined: set[str] = set() |
| 314 | for addr, callees in forward.items(): |
| 315 | if addr.split("::")[-1].split(".")[-1] == bare: |
| 316 | combined.update(callees) |
| 317 | initial = frozenset(combined) |
| 318 | |
| 319 | result: dict[int, list[str]] = {} |
| 320 | visited: set[str] = {start_addr} |
| 321 | q: deque[tuple[str, int]] = deque( |
| 322 | (name, 1) for name in sorted(initial) |
| 323 | ) |
| 324 | |
| 325 | while q: |
| 326 | name, depth = q.popleft() |
| 327 | if name in visited: |
| 328 | continue |
| 329 | visited.add(name) |
| 330 | result.setdefault(depth, []).append(name) |
| 331 | |
| 332 | if max_depth > 0 and depth >= max_depth: |
| 333 | continue |
| 334 | |
| 335 | # Find forward entries matching this bare name and recurse. |
| 336 | next_depth = depth + 1 |
| 337 | for addr, callees in forward.items(): |
| 338 | addr_bare = addr.split("::")[-1].split(".")[-1] |
| 339 | if addr_bare == name: |
| 340 | for callee in sorted(callees): |
| 341 | if callee not in visited: |
| 342 | q.append((callee, next_depth)) |
| 343 | |
| 344 | return result |
File History
1 commit
sha256:133f9bcf57a62ec7ebf0cd71138b54a200b15989aea2fb2a6912497e2926aa8a
feat(dev-safety): Phase 1 of #185 — audit tool for ambiguou…
Sonnet 5
patch
2 days ago