gabriel / muse public
_callgraph.py python
325 lines 11.2 KB
Raw
sha256:660fcac1df3ab28f61862e961890bd2ca8b754fa0242079d93ca1e25037ec8a6 chore(tests): add docstring to tests/__init__.py so rc14 tr… Human 26 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 from __future__ import annotations
32
33 import ast
34 import logging
35 import pathlib
36 from collections import deque
37
38 from muse.core.object_store import read_object
39 from muse.core.symbol_cache import SymbolCache
40 from muse.plugins.code.ast_parser import SymbolTree, parse_symbols
41 from muse.core.validation import MAX_AST_BYTES
42 from muse.core._types import Manifest
43 from muse.plugins.code._framework import ( # re-exported for caller convenience
44 ImplicitEntryEdge,
45 ImplicitEdgeGraph,
46 build_implicit_edge_graph,
47 )
48
49 logger = logging.getLogger(__name__)
50
51 #: Mapping from caller symbol address to the set of bare callee names.
52 ForwardGraph = dict[str, frozenset[str]]
53
54 #: Mapping from bare callee name to the list of caller addresses.
55 ReverseGraph = dict[str, list[str]]
56
57 _PY_SUFFIXES: frozenset[str] = frozenset({".py", ".pyi"})
58
59
60 # ---------------------------------------------------------------------------
61 # AST helpers
62 # ---------------------------------------------------------------------------
63
64
65 def call_name(func_node: ast.expr) -> str | None:
66 """Return the bare callee name from an ``ast.Call`` func node, or None.
67
68 Handles simple names (``foo()``) and attribute access (``obj.method()``).
69 Ignores subscript calls and other exotic forms.
70 """
71 if isinstance(func_node, ast.Name):
72 return func_node.id
73 if isinstance(func_node, ast.Attribute):
74 return func_node.attr
75 return None
76
77
78 def find_func_node(
79 stmts: list[ast.stmt],
80 name_parts: list[str],
81 ) -> ast.FunctionDef | ast.AsyncFunctionDef | None:
82 """Recursively locate a function node by its dotted qualified-name parts.
83
84 Args:
85 stmts: Statement list from a module or class body.
86 name_parts: Dotted path components, e.g. ``["User", "save"]``.
87
88 Returns:
89 The matching ``FunctionDef``/``AsyncFunctionDef`` node, or ``None``.
90 """
91 if not name_parts:
92 return None
93 target = name_parts[0]
94 for stmt in stmts:
95 if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)) and stmt.name == target:
96 if len(name_parts) == 1:
97 return stmt
98 elif isinstance(stmt, ast.ClassDef) and stmt.name == target and len(name_parts) > 1:
99 return find_func_node(stmt.body, name_parts[1:])
100 return None
101
102
103 def callees_for_symbol(source: bytes, address: str) -> list[str]:
104 """Return sorted unique bare callee names called by the Python symbol at *address*.
105
106 Args:
107 source: Raw bytes of the Python source file.
108 address: Full symbol address, e.g. ``"src/billing.py::compute_invoice_total"``.
109
110 Returns:
111 Sorted list of bare callee names. Empty if the file has a syntax
112 error, the symbol is not found, or it is not a function/method.
113 """
114 if "::" not in address:
115 return []
116 sym_qualified = address.split("::", 1)[1]
117 if len(source) > MAX_AST_BYTES:
118 return []
119 try:
120 tree = ast.parse(source)
121 except SyntaxError:
122 return []
123 func_node = find_func_node(tree.body, sym_qualified.split("."))
124 if func_node is None:
125 return []
126 names: set[str] = set()
127 for node in ast.walk(func_node):
128 if isinstance(node, ast.Call):
129 name = call_name(node.func)
130 if name:
131 names.add(name)
132 return sorted(names)
133
134
135 # ---------------------------------------------------------------------------
136 # Graph construction
137 # ---------------------------------------------------------------------------
138
139
140 def build_forward_graph(
141 root: pathlib.Path,
142 manifest: Manifest,
143 cache: SymbolCache | None = None,
144 ) -> ForwardGraph:
145 """Build the forward call graph from *manifest*.
146
147 Scans every Python file in the manifest, walking each symbol's body for
148 ``ast.Call`` nodes. Symbol trees are served from *cache* when provided,
149 avoiding redundant AST parses when the caller already has a warm cache.
150
151 Args:
152 root: Repository root (used to read blobs from the object store).
153 manifest: Snapshot manifest mapping file path → SHA-256 object ID.
154 cache: Optional shared ``SymbolCache`` instance. When ``None`` the
155 graph is built from raw bytes with no symbol caching.
156
157 Returns:
158 ``{caller_address: frozenset[callee_bare_name]}``.
159 """
160 graph: ForwardGraph = {}
161 for file_path, obj_id in manifest.items():
162 if pathlib.PurePosixPath(file_path).suffix.lower() not in _PY_SUFFIXES:
163 continue
164 raw = read_object(root, obj_id)
165 if raw is None:
166 continue
167 if len(raw) > MAX_AST_BYTES:
168 continue
169 try:
170 ast_tree = ast.parse(raw)
171 except SyntaxError:
172 continue
173 # Use the cache for symbol tree if available; fall back to direct parse.
174 sym_tree: SymbolTree = (
175 cache.get(obj_id) or parse_symbols(raw, file_path)
176 if cache is not None
177 else parse_symbols(raw, file_path)
178 )
179 for addr, rec in sym_tree.items():
180 if rec["kind"] not in {"function", "async_function", "method", "async_method"}:
181 continue
182 func_node = find_func_node(ast_tree.body, rec["qualified_name"].split("."))
183 if func_node is None:
184 continue
185 names: set[str] = set()
186 for node in ast.walk(func_node):
187 if isinstance(node, ast.Call):
188 name = call_name(node.func)
189 if name:
190 names.add(name)
191 graph[addr] = frozenset(names)
192 return graph
193
194
195 def build_reverse_graph(
196 root: pathlib.Path,
197 manifest: Manifest,
198 cache: SymbolCache | None = None,
199 ) -> ReverseGraph:
200 """Build the reverse call graph from *manifest*.
201
202 Inverts the forward graph: maps each callee bare name to every caller
203 address that calls it.
204
205 Args:
206 root: Repository root.
207 manifest: Snapshot manifest mapping file path → SHA-256 object ID.
208 cache: Optional shared ``SymbolCache`` instance passed through to
209 :func:`build_forward_graph`.
210
211 Returns:
212 ``{callee_bare_name: [caller_address, ...]}``.
213 """
214 forward = build_forward_graph(root, manifest, cache=cache)
215 reverse: ReverseGraph = {}
216 for caller_addr, callee_names in forward.items():
217 for name in callee_names:
218 reverse.setdefault(name, []).append(caller_addr)
219 # Sort each caller list for deterministic output.
220 for name in reverse:
221 reverse[name].sort()
222 return reverse
223
224
225 def transitive_callers(
226 start_name: str,
227 reverse: ReverseGraph,
228 max_depth: int = 0,
229 ) -> dict[int, list[str]]:
230 """BFS through *reverse* to find all transitive callers of *start_name*.
231
232 Args:
233 start_name: Bare function/method name to start from.
234 reverse: Reverse call graph produced by :func:`build_reverse_graph`.
235 max_depth: Maximum BFS depth. ``0`` means unlimited.
236
237 Returns:
238 ``{depth: [caller_address, ...]}``, depth 1 = direct callers.
239 Addresses that appear at multiple depths are recorded only at the
240 shallowest depth (first encounter wins).
241 """
242 result: dict[int, list[str]] = {}
243 visited_names: set[str] = {start_name}
244 visited_addrs: set[str] = set()
245 # Use deque for O(1) popleft instead of O(n) list.pop(0).
246 q: deque[tuple[str, int]] = deque([(start_name, 0)])
247
248 while q:
249 name, depth = q.popleft()
250 if max_depth > 0 and depth >= max_depth:
251 continue
252 next_depth = depth + 1
253 for caller_addr in reverse.get(name, []):
254 if caller_addr in visited_addrs:
255 continue
256 visited_addrs.add(caller_addr)
257 result.setdefault(next_depth, []).append(caller_addr)
258 # Extract the bare name of the caller to continue BFS.
259 caller_name = caller_addr.split("::")[-1].split(".")[-1]
260 if caller_name not in visited_names:
261 visited_names.add(caller_name)
262 q.append((caller_name, next_depth))
263
264 return result
265
266
267 def transitive_callees(
268 start_addr: str,
269 forward: ForwardGraph,
270 max_depth: int = 0,
271 ) -> dict[int, list[str]]:
272 """BFS through *forward* to find all transitive callees of *start_addr*.
273
274 Args:
275 start_addr: Full caller address (``file::symbol``) or bare name to
276 start from.
277 forward: Forward call graph produced by :func:`build_forward_graph`.
278 max_depth: Maximum BFS depth. ``0`` means unlimited.
279
280 Returns:
281 ``{depth: [callee_bare_name, ...]}``, depth 1 = direct callees.
282 Names that appear at multiple depths are recorded only at the
283 shallowest depth (first encounter wins).
284 """
285 # Seed from the start address's own callee set if it exists; otherwise
286 # treat start_addr as a bare name and seed from any forward entry whose
287 # address ends with that name.
288 initial: frozenset[str]
289 if start_addr in forward:
290 initial = forward[start_addr]
291 else:
292 # Bare-name fallback: collect callees from any matching address.
293 bare = start_addr.split("::")[-1].split(".")[-1]
294 combined: set[str] = set()
295 for addr, callees in forward.items():
296 if addr.split("::")[-1].split(".")[-1] == bare:
297 combined.update(callees)
298 initial = frozenset(combined)
299
300 result: dict[int, list[str]] = {}
301 visited: set[str] = {start_addr}
302 q: deque[tuple[str, int]] = deque(
303 (name, 1) for name in sorted(initial)
304 )
305
306 while q:
307 name, depth = q.popleft()
308 if name in visited:
309 continue
310 visited.add(name)
311 result.setdefault(depth, []).append(name)
312
313 if max_depth > 0 and depth >= max_depth:
314 continue
315
316 # Find forward entries matching this bare name and recurse.
317 next_depth = depth + 1
318 for addr, callees in forward.items():
319 addr_bare = addr.split("::")[-1].split(".")[-1]
320 if addr_bare == name:
321 for callee in sorted(callees):
322 if callee not in visited:
323 q.append((callee, next_depth))
324
325 return result
File History 5 commits
sha256:660fcac1df3ab28f61862e961890bd2ca8b754fa0242079d93ca1e25037ec8a6 chore(tests): add docstring to tests/__init__.py so rc14 tr… Human 26 days ago
sha256:d8316ffae901be06347e16ab55be11868eb519dd16ade3e8aa16a99e662f7e62 baseline: rc14 re-baseline after rc3 store corruption recovery Human patch 26 days ago
sha256:50d413a635f57761c3fce1f70251b957b6423821525bf330747ef4007339222e Merge branch 'dev' into main Human 29 days ago
sha256:fb67fed5a4d3e40de84bdd163de94ef1386570bef1dd1a020a732c8a038962ce Merge branch 'dev' into main Human 48 days ago
sha256:1c4b3e3a9a1f300774c3ee662b572a698d5fd405bf765a71e6011a2e9c3eaaaa feat: Muse — version control for the agent era Human 100 days ago