"""muse code gravity — structural weight of every symbol. "Gravity" answers the question every architect dreads: **"If I change the contract of this symbol, how much of the codebase breaks?"** Not just the direct callers — the full transitive closure. If ``read_object`` changes, that breaks ``read_commit``, which breaks ``resolve_commit_ref``, which breaks every route handler, which breaks every CLI command that touches the repo. That is gravity: the fraction of the production codebase that lives downstream of a symbol in the call graph. Why this matters ---------------- * High-gravity symbols are **structural pillars**. Treat their public contracts as APIs — design changes carefully, communicate broadly, bump MAJOR. * Low-gravity symbols are **safe to refactor**. They sit at the edge of the dependency graph with few downstream effects. * **Gravity ≠ hotspot.** A frequently-changed symbol with low gravity is a quick iteration loop. A rarely-changed symbol with gravity 90% is a time bomb — the moment its contract shifts, everything breaks. How it works ------------ 1. Loads the committed HEAD snapshot and builds the production call graph (AST-based, same approach as ``blast-risk`` and ``semantic-test-coverage``). 2. Inverts the graph: for each symbol, BFS through the reverse call graph to find every symbol that transitively depends on it. 3. ``gravity_pct = transitive_dependents / total_production_symbols × 100``. Test files are excluded from the dependency graph by default (a function being called by 50 test functions does not mean it has production structural weight). Pass ``--include-tests`` to count test callers as well. Usage:: muse code gravity muse code gravity --top 20 muse code gravity --min-gravity 50 muse code gravity --kind function muse code gravity --file core/store.py muse code gravity --sort depth muse code gravity --depth 3 muse code gravity --explain "muse/core/object_store.py::read_object" muse code gravity --include-tests muse code gravity --json Default output:: Symbol gravity — HEAD (378 production symbols) Structural weight: fraction of production codebase that transitively depends on each symbol. # ADDRESS GRAVITY DIRECT DEPTH 1 muse/core/object_store.py::read_object 94.2% 28 8 2 muse/core/store.py::resolve_commit_ref 87.1% 47 6 3 muse/core/errors.py::ExitCode 75.4% 31 5 ⚠️ High-gravity symbols are structural pillars. Treat their contracts as APIs. ``--explain`` output:: Gravity breakdown — muse/core/object_store.py::read_object Kind: function Total gravity: 356 / 378 (94.2%) Direct callers: 28 Max depth: 8 Depth distribution: depth 1 (direct): 28 callers ████████████████████ depth 2: 89 callers ████████████████████████████████████████ depth 3: 127 callers ████████████████████████████████████████████████████ depth 4: 82 callers ██████████████████████████████████ depth 5+: 30 callers ████████████ Deepest callers (depth 8): muse/api/routes/wire.py::push_handler muse/api/routes/wire.py::pull_handler muse/cli/commands/log.py::run A breaking change here propagates through 8 layers of the codebase. JSON output (``--json``):: { "ref": "HEAD", "snapshot_id": "a3f2c9e1ab2c", "total_production_symbols": 378, "max_depth": 0, "include_tests": false, "filters": { "kind": null, "file": null, "min_gravity": 0.0, "top": 20 }, "symbols": [ { "address": "muse/core/object_store.py::read_object", "name": "read_object", "kind": "function", "file": "muse/core/object_store.py", "gravity_pct": 94.2, "direct_dependents": 28, "transitive_dependents": 356, "max_depth": 8, "depth_distribution": { "1": 28, "2": 89, "3": 127, "4": 82, "5": 30 } } ] } Security note ------------- Symbol addresses are validated to contain ``::`` before any store access. File paths from the manifest are not passed to the shell — all access goes through the object store's content-addressed lookup. """ from __future__ import annotations import argparse import ast import json import logging import pathlib import re import sys from typing import Literal, TypedDict from muse.core.errors import ExitCode from muse.core.repo import read_repo_id, require_repo from muse.core.store import ( Manifest, get_commit_snapshot_manifest, read_current_branch, resolve_commit_ref, ) from muse.plugins.code._callgraph import ( ForwardGraph, ReverseGraph, call_name, find_func_node, transitive_callers, ) from muse.plugins.code._query import symbols_for_snapshot from muse.plugins.code.ast_parser import SymbolRecord, parse_symbols from muse.core.validation import clamp_int, MAX_AST_BYTES, sanitize_display type _BlobMap = dict[str, bytes] type _SymbolIndex = dict[str, SymbolRecord] type _CounterMap = dict[str, int] logger = logging.getLogger(__name__) # ── Constants ────────────────────────────────────────────────────────────────── _DEFAULT_TOP = 20 _DEFAULT_DEPTH = 0 # 0 = unlimited _PY_SUFFIXES: frozenset[str] = frozenset({".py", ".pyi"}) _TEST_PATTERNS: tuple[re.Pattern[str], ...] = ( re.compile(r"(^|/)test_"), re.compile(r"_test\.py$"), re.compile(r"(^|/)tests/"), re.compile(r"(^|/)spec/"), re.compile(r"(^|/)conftest\.py$"), ) _IMPORT_KIND = "import" _TRACKED_KINDS: frozenset[str] = frozenset( {"function", "async_function", "method", "async_method", "class"} ) SortKey = Literal["gravity", "direct", "depth"] _SORT_CHOICES: tuple[SortKey, ...] = ("gravity", "direct", "depth") _BAR_MAX_WIDTH = 40 # ── TypedDicts ───────────────────────────────────────────────────────────────── class _SymbolGravity(TypedDict): """Gravity record for a single production symbol.""" address: str name: str kind: str file: str gravity_pct: float direct_dependents: int transitive_dependents: int max_depth: int depth_distribution: _CounterMap class _FilterSpec(TypedDict): kind: str | None file: str | None min_gravity: float top: int class _JsonOut(TypedDict): ref: str snapshot_id: str total_production_symbols: int max_depth: int include_tests: bool filters: _FilterSpec symbols: list[_SymbolGravity] # ── Helpers ──────────────────────────────────────────────────────────────────── def _is_test_file(file_path: str) -> bool: """Return True if *file_path* looks like a test or conftest file.""" return any(p.search(file_path) for p in _TEST_PATTERNS) def _load_py_blobs( root: pathlib.Path, manifest: Manifest, include_tests: bool, ) -> _BlobMap: """Load Python file blobs from the object store into memory. Args: root: Repository root. manifest: Snapshot manifest. include_tests: When False (default), test files are excluded. Returns: Mapping ``{file_path: raw_bytes}`` for qualifying Python files. """ from muse.core.object_store import read_object blobs: _BlobMap = {} for file_path, obj_id in manifest.items(): if pathlib.PurePosixPath(file_path).suffix.lower() not in _PY_SUFFIXES: continue if not include_tests and _is_test_file(file_path): continue raw = read_object(root, obj_id) if raw is not None: blobs[file_path] = raw return blobs def _build_forward_graph( blobs: _BlobMap, ) -> ForwardGraph: """Build the forward call graph from pre-loaded blobs. Scans every Python file in *blobs*, walking each callable symbol's body for ``ast.Call`` nodes. Returns ``{caller_address: frozenset[callee_bare_name]}``. Building from pre-loaded blobs avoids a second object-store round-trip. """ graph: ForwardGraph = {} for file_path, raw in blobs.items(): try: if len(raw) > MAX_AST_BYTES: return {} tree = ast.parse(raw) except SyntaxError: continue sym_tree = parse_symbols(raw, file_path) for addr, rec in sym_tree.items(): if rec["kind"] not in {"function", "async_function", "method", "async_method"}: continue func_node = find_func_node(tree.body, rec["qualified_name"].split(".")) if func_node is None: continue callees: set[str] = set() for node in ast.walk(func_node): if isinstance(node, ast.Call): n = call_name(node.func) if n: callees.add(n) graph[addr] = frozenset(callees) return graph def _invert_graph(forward: ForwardGraph) -> ReverseGraph: """Invert the forward call graph to get ``{callee_bare_name: [caller_addr, ...]}``. Produces a sorted list of callers for each callee name so output is deterministic across runs. """ reverse: ReverseGraph = {} for caller_addr, callee_names in forward.items(): for name in callee_names: reverse.setdefault(name, []).append(caller_addr) for name in reverse: reverse[name].sort() return reverse def _gravity_bar(count: int, max_count: int, width: int = _BAR_MAX_WIDTH) -> str: """Return a unicode block bar proportional to *count* / *max_count*.""" if max_count <= 0: return "" filled = round(count / max_count * width) return "█" * filled # ── Core algorithm ───────────────────────────────────────────────────────────── def _compute_gravity_for_symbol( bare_name: str, reverse: ReverseGraph, max_depth: int, ) -> tuple[dict[int, list[str]], int]: """Compute transitive dependents for a symbol with the given bare name. Uses the existing ``transitive_callers`` BFS from ``_callgraph``. Args: bare_name: Bare callee name (last component of the symbol address). reverse: Reverse call graph from ``_invert_graph``. max_depth: BFS depth limit; ``0`` = unlimited. Returns: ``(depth_map, max_reached_depth)`` where depth_map is ``{depth: [caller_address, ...]}``. """ depth_map = transitive_callers(bare_name, reverse, max_depth=max_depth) max_reached = max(depth_map) if depth_map else 0 return depth_map, max_reached def _build_gravity_records( prod_symbols: _SymbolIndex, reverse: ReverseGraph, max_depth: int, file_filter: str | None, kind_filter: str | None, min_gravity: float, top: int, sort_key: SortKey, ) -> tuple[list[_SymbolGravity], int]: """Compute gravity for every qualifying production symbol. Args: prod_symbols: Flat ``{address: SymbolRecord}`` for all production symbols. reverse: Reverse call graph. max_depth: BFS depth cap (``0`` = unlimited). file_filter: Optional path suffix filter. kind_filter: Optional symbol kind filter. min_gravity: Minimum gravity percentage threshold. top: Maximum number of records to return (``0`` = unlimited). sort_key: Sort dimension. Returns: ``(records, total_prod_symbols)`` — records are sorted and truncated to *top*. ``total_prod_symbols`` is the unfiltered production symbol count used as the denominator for gravity %. """ total = len(prod_symbols) if total == 0: return [], 0 records: list[_SymbolGravity] = [] for addr, rec in prod_symbols.items(): kind = rec["kind"] if kind == _IMPORT_KIND or kind not in _TRACKED_KINDS: continue file_path = addr.split("::")[0] if file_filter and file_filter not in file_path: continue if kind_filter and kind != kind_filter: continue bare_name = rec["name"] depth_map, max_reached = _compute_gravity_for_symbol(bare_name, reverse, max_depth) direct = len(depth_map.get(1, [])) total_trans = sum(len(v) for v in depth_map.values()) denom = max(1, total - 1) # exclude self pct = round(total_trans / denom * 100, 1) if pct < min_gravity: continue # Flatten depth distribution to str-keyed dict for JSON. dist: _CounterMap = {str(d): len(addrs) for d, addrs in depth_map.items()} records.append( _SymbolGravity( address=addr, name=bare_name, kind=kind, file=file_path, gravity_pct=pct, direct_dependents=direct, transitive_dependents=total_trans, max_depth=max_reached, depth_distribution=dist, ) ) # Sort. if sort_key == "direct": records.sort(key=lambda r: (r["direct_dependents"], r["gravity_pct"]), reverse=True) elif sort_key == "depth": records.sort(key=lambda r: (r["max_depth"], r["gravity_pct"]), reverse=True) else: # gravity (default) records.sort(key=lambda r: (r["gravity_pct"], r["transitive_dependents"]), reverse=True) if top > 0: records = records[:top] return records, total def _find_deepest_callers( bare_name: str, reverse: ReverseGraph, max_depth: int, show_count: int = 5, ) -> tuple[list[str], int]: """Return the addresses of the deepest callers (longest dependency chains). Args: bare_name: Target symbol's bare name. reverse: Reverse call graph. max_depth: BFS depth cap; 0 = unlimited. show_count: Maximum number of deepest-caller addresses to return. Returns: ``(addresses, depth)`` — the deepest addresses found, and the depth at which they were found. """ depth_map, max_reached = _compute_gravity_for_symbol(bare_name, reverse, max_depth) if not depth_map: return [], 0 deepest = depth_map[max_reached] return deepest[:show_count], max_reached # ── Formatters ───────────────────────────────────────────────────────────────── def _print_leaderboard( records: list[_SymbolGravity], total_prod: int, max_depth_cap: int, sort_key: SortKey, include_tests: bool, ) -> None: """Print the human-readable gravity leaderboard.""" scope = "all symbols" if include_tests else "production symbols" cap_str = f" · depth cap: {max_depth_cap}" if max_depth_cap > 0 else "" print(f"\n Symbol gravity — HEAD ({total_prod} {scope}{cap_str})") print( f" Structural weight: fraction of {scope} that transitively depends" " on each symbol.\n" ) if not records: print(" No symbols match the current filters.\n") return # Column widths. addr_w = min(60, max(len(r["address"]) for r in records) + 2) print( f" {'#':>4} {'ADDRESS':<{addr_w}} {'GRAVITY':>8} {'DIRECT':>7} {'DEPTH':>6}" ) print(f" {'─' * 4} {'─' * addr_w} {'─' * 8} {'─' * 7} {'─' * 6}") for i, rec in enumerate(records, 1): addr = rec["address"] if len(addr) > addr_w: addr = "…" + addr[-(addr_w - 1):] pct_str = f"{rec['gravity_pct']:5.1f}%" print( f" {i:>4} {addr:<{addr_w}} {pct_str:>8}" f" {rec['direct_dependents']:>7}" f" {rec['max_depth']:>6}" ) print() if records: top = records[0] print( f" ⚠️ {top['name']} carries {top['gravity_pct']:.1f}% structural weight" " — treat its contract as an API." ) print() def _print_explain( rec: _SymbolGravity, reverse: ReverseGraph, max_depth: int, total_prod: int, ) -> None: """Print the detailed gravity breakdown for a single symbol.""" print(f"\n Gravity breakdown — {sanitize_display(rec['address'])}\n") print(f" Kind: {sanitize_display(rec['kind'])}") print(f" File: {sanitize_display(rec['file'])}") print( f" Total gravity: {rec['transitive_dependents']} / {total_prod}" f" ({rec['gravity_pct']:.1f}%)" ) print(f" Direct callers: {rec['direct_dependents']}") print(f" Max depth: {rec['max_depth']}") print() if rec["depth_distribution"]: # Re-fetch depth_map for the distribution bars. depth_map, _ = _compute_gravity_for_symbol(rec["name"], reverse, max_depth) max_at_level = max(len(v) for v in depth_map.values()) if depth_map else 1 print(" Depth distribution:") for depth in sorted(depth_map): callers = depth_map[depth] count = len(callers) bar = _gravity_bar(count, max_at_level, width=30) label = f"depth {depth} (direct):" if depth == 1 else f"depth {depth}: " print(f" {label} {count:>5} callers {bar}") print() # Deepest callers. deepest_addrs, deepest_depth = _find_deepest_callers(rec["name"], reverse, max_depth) if deepest_addrs: print(f" Deepest callers (depth {deepest_depth}):") for addr in deepest_addrs: print(f" {sanitize_display(addr)}") print() print( f" A breaking change here propagates through" f" {deepest_depth} layers of the codebase." ) print() # ── Entry point ──────────────────────────────────────────────────────────────── def run(args: argparse.Namespace) -> None: """Entry point for ``muse code gravity``.""" root = require_repo() from muse.plugins.registry import read_domain try: if read_domain(root) == "knowtation": from muse.plugins.knowtation.cli_hooks import run_gravity_for_vault run_gravity_for_vault(root, args) return except Exception: # noqa: BLE001 pass # ── Argument validation ──────────────────────────────────────────────────── top: int = clamp_int(args.top, 0, 10000, 'top') if top < 0: print("❌ --top must be >= 0 (0 = unlimited).", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) max_depth: int = clamp_int(args.depth, 0, 50, 'depth') if max_depth < 0: print("❌ --depth must be >= 0 (0 = unlimited).", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) min_gravity: float = args.min_gravity if not (0.0 <= min_gravity <= 100.0): print("❌ --min-gravity must be between 0.0 and 100.0.", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) sort_key: SortKey = args.sort kind_filter: str | None = args.kind or None file_filter: str | None = args.file or None include_tests: bool = args.include_tests explain_addr: str | None = args.explain or None if explain_addr is not None and "::" not in explain_addr: print( "❌ --explain ADDRESS must be in ``file::Symbol`` format.", file=sys.stderr, ) raise SystemExit(ExitCode.USER_ERROR) # ── Resolve HEAD ─────────────────────────────────────────────────────────── repo_id = read_repo_id(root) branch = read_current_branch(root) head = resolve_commit_ref(root, repo_id, branch, None) if head is None: print("❌ HEAD commit not found — is this an empty repository?", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) manifest = get_commit_snapshot_manifest(root, head.commit_id) or {} # ── Single bulk read of all qualifying Python blobs ──────────────────────── blobs = _load_py_blobs(root, manifest, include_tests) # ── Symbol extraction ────────────────────────────────────────────────────── all_trees = symbols_for_snapshot(root, manifest) # Partition: test files vs. production (for symbol scope). prod_trees = ( all_trees if include_tests else {fp: t for fp, t in all_trees.items() if not _is_test_file(fp)} ) # Flat production symbol index (imports excluded). prod_symbols: _SymbolIndex = {} for sym_tree in prod_trees.values(): for addr, rec in sym_tree.items(): if rec["kind"] != _IMPORT_KIND and rec["kind"] in _TRACKED_KINDS: prod_symbols[addr] = rec if not prod_symbols: print(" No production symbols found in HEAD snapshot.\n") return # ── Build call graphs ────────────────────────────────────────────────────── forward = _build_forward_graph(blobs) reverse = _invert_graph(forward) # ── --explain: single-symbol deep dive ──────────────────────────────────── if explain_addr is not None: if explain_addr not in prod_symbols: print( f"❌ {explain_addr!r} not found in HEAD production symbols.", file=sys.stderr, ) print( " Use ``muse code symbols`` to list available addresses.", file=sys.stderr, ) raise SystemExit(ExitCode.USER_ERROR) rec_explain = prod_symbols[explain_addr] bare = rec_explain["name"] depth_map, max_reached = _compute_gravity_for_symbol(bare, reverse, max_depth) direct = len(depth_map.get(1, [])) total_trans = sum(len(v) for v in depth_map.values()) denom = max(1, len(prod_symbols) - 1) pct = round(total_trans / denom * 100, 1) dist = {str(d): len(addrs) for d, addrs in depth_map.items()} single_rec = _SymbolGravity( address=explain_addr, name=bare, kind=rec_explain["kind"], file=explain_addr.split("::")[0], gravity_pct=pct, direct_dependents=direct, transitive_dependents=total_trans, max_depth=max_reached, depth_distribution=dist, ) if args.json: print(json.dumps(single_rec, indent=2)) return _print_explain(single_rec, reverse, max_depth, len(prod_symbols)) return # ── Leaderboard ──────────────────────────────────────────────────────────── records, total_prod = _build_gravity_records( prod_symbols, reverse, max_depth, file_filter, kind_filter, min_gravity, top, sort_key, ) if args.json: out: _JsonOut = _JsonOut( ref="HEAD", snapshot_id=head.commit_id[:12], total_production_symbols=total_prod, max_depth=max_depth, include_tests=include_tests, filters=_FilterSpec( kind=kind_filter, file=file_filter, min_gravity=min_gravity, top=top, ), symbols=records, ) print(json.dumps(out, indent=2)) return _print_leaderboard(records, total_prod, max_depth, sort_key, include_tests) # ── CLI registration ─────────────────────────────────────────────────────────── def register( sub: argparse._SubParsersAction[argparse.ArgumentParser], ) -> None: """Register ``gravity`` under the ``code`` subcommand group. Args: sub: The subparser action from the ``code`` command group. """ p = sub.add_parser( "gravity", help=( "Structural weight — how much of the production codebase" " transitively depends on each symbol." ), description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) p.add_argument( "--explain", metavar="ADDRESS", help=( "Deep-dive into a specific symbol's gravity:" " depth distribution and deepest callers." ), ) p.add_argument( "--top", type=int, default=_DEFAULT_TOP, metavar="N", help=f"Show only the top N symbols by gravity (default: {_DEFAULT_TOP}; 0 = all).", ) p.add_argument( "--sort", metavar="KEY", choices=list(_SORT_CHOICES), default="gravity", help=( "Sort dimension: ``gravity`` (default, transitive %%)" ", ``direct`` (direct callers)" ", ``depth`` (max dependency chain length)." ), ) p.add_argument( "--depth", type=int, default=_DEFAULT_DEPTH, metavar="D", help=( f"Maximum BFS depth for transitive analysis (default: {_DEFAULT_DEPTH} = unlimited)." " Use a small value (e.g. 3) for large repos to bound runtime." ), ) p.add_argument( "--kind", metavar="KIND", choices=["function", "async_function", "method", "async_method", "class"], help="Filter symbols by kind.", ) p.add_argument( "--file", metavar="SUFFIX", help="Scope analysis to symbols in files whose path contains SUFFIX.", ) p.add_argument( "--min-gravity", type=float, default=0.0, metavar="PCT", help="Show only symbols with gravity >= PCT%%.", ) p.add_argument( "--include-tests", action="store_true", help=( "Count test-file callers in the dependency graph" " (default: test callers are excluded)." ), ) p.add_argument( "--json", action="store_true", help="Emit JSON instead of human-readable text.", ) p.set_defaults(func=run)