_query.py
file-level
1
files
1
commits
0
hotspots
0
🧊 dead
0
💥 blast risk
| 1 | """Shared query helpers for the code-domain CLI commands. |
| 2 | |
| 3 | This module provides the low-level primitives that multiple code-domain |
| 4 | commands need — symbol extraction from snapshots, commit-graph walking, |
| 5 | and language classification — so each command can stay thin. |
| 6 | |
| 7 | None of these functions are part of the public ``CodePlugin`` API. They |
| 8 | are internal helpers for the CLI layer and must not be imported by any |
| 9 | core module. |
| 10 | """ |
| 11 | |
| 12 | from __future__ import annotations |
| 13 | |
| 14 | import hashlib |
| 15 | import itertools |
| 16 | import logging |
| 17 | import pathlib |
| 18 | from collections import deque |
| 19 | from collections.abc import Iterator |
| 20 | |
| 21 | from muse.core._types import Manifest, Metadata |
| 22 | from muse.core.object_store import read_object |
| 23 | from muse.core.store import CommitRecord, read_commit |
| 24 | from muse.core.symbol_cache import SymbolCache, load_symbol_cache |
| 25 | from muse.domain import DomainOp |
| 26 | from muse.plugins.code.ast_parser import ( |
| 27 | SymbolTree, |
| 28 | parse_symbols, |
| 29 | ) |
| 30 | |
| 31 | type _SymbolTreeMap = dict[str, SymbolTree] |
| 32 | type _CommitIndex = dict[str, CommitRecord] |
| 33 | |
| 34 | logger = logging.getLogger(__name__) |
| 35 | |
| 36 | # --------------------------------------------------------------------------- |
| 37 | # Language classification |
| 38 | # --------------------------------------------------------------------------- |
| 39 | |
| 40 | _SUFFIX_LANG: Metadata = { |
| 41 | ".py": "Python", ".pyi": "Python", |
| 42 | ".ts": "TypeScript", ".tsx": "TypeScript", |
| 43 | ".js": "JavaScript", ".jsx": "JavaScript", |
| 44 | ".mjs": "JavaScript", ".cjs": "JavaScript", |
| 45 | ".go": "Go", |
| 46 | ".rs": "Rust", |
| 47 | ".java": "Java", |
| 48 | ".cs": "C#", |
| 49 | ".c": "C", ".h": "C", |
| 50 | ".cpp": "C++", ".cc": "C++", ".cxx": "C++", ".hpp": "C++", ".hxx": "C++", |
| 51 | ".rb": "Ruby", |
| 52 | ".kt": "Kotlin", ".kts": "Kotlin", |
| 53 | ".swift": "Swift", |
| 54 | ".md": "Markdown", ".rst": "reStructuredText", ".txt": "Text", |
| 55 | ".toml": "TOML", |
| 56 | ".yaml": "YAML", ".yml": "YAML", |
| 57 | ".json": "JSON", ".jsonc": "JSON", |
| 58 | ".css": "CSS", ".scss": "SCSS", |
| 59 | ".html": "HTML", ".htm": "HTML", |
| 60 | ".sql": "SQL", |
| 61 | ".sh": "Shell", ".bash": "Shell", ".zsh": "Shell", |
| 62 | ".proto": "Protobuf", |
| 63 | ".tf": "Terraform", |
| 64 | } |
| 65 | |
| 66 | |
| 67 | def language_of(file_path: str) -> str: |
| 68 | """Return a display language name for *file_path* based on its suffix.""" |
| 69 | suffix = pathlib.PurePosixPath(file_path).suffix.lower() |
| 70 | return _SUFFIX_LANG.get(suffix, suffix or "(no ext)") |
| 71 | |
| 72 | |
| 73 | def is_semantic(file_path: str) -> bool: |
| 74 | """Return ``True`` if *file_path* has a suffix with AST-level support.""" |
| 75 | from muse.plugins.code.ast_parser import SEMANTIC_EXTENSIONS |
| 76 | suffix = pathlib.PurePosixPath(file_path).suffix.lower() |
| 77 | return suffix in SEMANTIC_EXTENSIONS |
| 78 | |
| 79 | |
| 80 | # --------------------------------------------------------------------------- |
| 81 | # Symbol extraction from a snapshot manifest |
| 82 | # --------------------------------------------------------------------------- |
| 83 | |
| 84 | |
| 85 | def _rekey_tree(tree: SymbolTree, file_path: str) -> SymbolTree: |
| 86 | """Re-key *tree* so every address starts with *file_path*. |
| 87 | |
| 88 | The ``SymbolCache`` is keyed by content SHA-256. When two files have |
| 89 | identical bytes they share one cache entry whose addresses carry the path |
| 90 | of whichever file was parsed first. This function patches those addresses |
| 91 | to the actual *file_path* being processed — an O(n) pass that is only |
| 92 | triggered on a cache hit where the prefix differs. |
| 93 | """ |
| 94 | first = next(iter(tree), None) |
| 95 | if first is None: |
| 96 | return tree |
| 97 | cached_prefix = first.split("::")[0] if "::" in first else None |
| 98 | if cached_prefix == file_path: |
| 99 | return tree # Addresses already correct — common case. |
| 100 | return { |
| 101 | f"{file_path}::{addr.split('::', 1)[1]}" if "::" in addr else addr: rec |
| 102 | for addr, rec in tree.items() |
| 103 | } |
| 104 | |
| 105 | |
| 106 | def symbols_for_snapshot( |
| 107 | root: pathlib.Path, |
| 108 | manifest: Manifest, |
| 109 | *, |
| 110 | kind_filter: str | None = None, |
| 111 | file_filter: str | None = None, |
| 112 | language_filter: str | None = None, |
| 113 | workdir: pathlib.Path | None = None, |
| 114 | cache: SymbolCache | None = None, |
| 115 | ) -> _SymbolTreeMap: |
| 116 | """Extract symbol trees for all semantic files in *manifest*. |
| 117 | |
| 118 | Results are served from the persistent symbol cache when available, |
| 119 | cutting a full-snapshot scan from ~1,300 ms to ~22 ms on a warm cache. |
| 120 | The cache is keyed by the SHA-256 of the file bytes (``object_id`` for |
| 121 | committed files; freshly computed SHA-256 for working-tree reads), so |
| 122 | every hit is guaranteed correct. |
| 123 | |
| 124 | Args: |
| 125 | root: Repository root (used to locate the object store and |
| 126 | load/save the symbol cache). |
| 127 | manifest: Snapshot manifest mapping file path → SHA-256. |
| 128 | kind_filter: If set, only include symbols with this ``kind``. |
| 129 | file_filter: If set, only include symbols from this exact file path. |
| 130 | language_filter: If set, only include symbols from files of this language. |
| 131 | workdir: When set, file content is read from *workdir* first |
| 132 | (working-tree mode), falling back to the object store if |
| 133 | the file does not exist on disk. Pass ``root`` to get |
| 134 | live, uncommitted changes reflected immediately. |
| 135 | cache: Optional pre-loaded ``SymbolCache`` instance. When |
| 136 | ``None`` the cache is loaded from disk automatically |
| 137 | and saved at the end of the call. Pass an explicit |
| 138 | instance when batching multiple calls to share one |
| 139 | cache load/save cycle. |
| 140 | |
| 141 | Returns: |
| 142 | Dict mapping ``file_path → SymbolTree``; empty trees are omitted. |
| 143 | """ |
| 144 | if cache is None: |
| 145 | active_cache: SymbolCache = load_symbol_cache(root) |
| 146 | own_cache = True |
| 147 | else: |
| 148 | active_cache = cache |
| 149 | own_cache = False |
| 150 | |
| 151 | result: _SymbolTreeMap = {} |
| 152 | for file_path, object_id in sorted(manifest.items()): |
| 153 | if not is_semantic(file_path): |
| 154 | continue |
| 155 | if file_filter and file_path != file_filter: |
| 156 | continue |
| 157 | if language_filter and language_of(file_path) != language_filter: |
| 158 | continue |
| 159 | |
| 160 | # --- Determine cache key and read bytes if needed ---------------- |
| 161 | # For working-tree files the object_id is the committed version; |
| 162 | # the disk content may differ. We compute the SHA-256 of whatever |
| 163 | # bytes we actually parse so the cache key is always content-addressed. |
| 164 | cache_key = object_id |
| 165 | raw: bytes | None = None |
| 166 | |
| 167 | if workdir is not None: |
| 168 | disk_path = workdir / file_path |
| 169 | if disk_path.is_file(): |
| 170 | try: |
| 171 | raw = disk_path.read_bytes() |
| 172 | # Recompute the key for the actual disk content. |
| 173 | cache_key = hashlib.sha256(raw).hexdigest() |
| 174 | except OSError as exc: |
| 175 | logger.debug("Could not read %s from workdir: %s", file_path, exc) |
| 176 | |
| 177 | # --- Cache lookup ------------------------------------------------ |
| 178 | tree = active_cache.get(cache_key) |
| 179 | if tree is None: |
| 180 | # Cache miss — fetch bytes and parse. |
| 181 | if raw is None: |
| 182 | raw = read_object(root, object_id) |
| 183 | if raw is None: |
| 184 | logger.debug("Object %s missing for %s — skipping", object_id[:8], file_path) |
| 185 | continue |
| 186 | tree = parse_symbols(raw, file_path) |
| 187 | active_cache.put(cache_key, tree) |
| 188 | else: |
| 189 | # Cache hit — re-key addresses when two files share the same |
| 190 | # SHA-256 (identical content). The cache stores the tree built |
| 191 | # for whichever file was parsed first; its addresses carry that |
| 192 | # file's path prefix. Patch them to the current file_path so |
| 193 | # cross-file clone detection and all downstream consumers get |
| 194 | # correct symbol addresses. |
| 195 | tree = _rekey_tree(tree, file_path) |
| 196 | |
| 197 | # --- Apply filters on top of the (potentially cached) full tree -- |
| 198 | if kind_filter: |
| 199 | tree = {addr: rec for addr, rec in tree.items() if rec["kind"] == kind_filter} |
| 200 | if tree: |
| 201 | result[file_path] = tree |
| 202 | |
| 203 | if own_cache: |
| 204 | active_cache.save() |
| 205 | |
| 206 | return result |
| 207 | |
| 208 | |
| 209 | # --------------------------------------------------------------------------- |
| 210 | # Commit-graph walking |
| 211 | # --------------------------------------------------------------------------- |
| 212 | |
| 213 | |
| 214 | def walk_commits_bfs( |
| 215 | root: pathlib.Path, |
| 216 | start_commit_id: str, |
| 217 | max_commits: int = 500, |
| 218 | stop_at_commit_id: str | None = None, |
| 219 | ) -> tuple[list[CommitRecord], bool]: |
| 220 | """BFS walk of the commit DAG from *start_commit_id*, newest-first. |
| 221 | |
| 222 | Unlike the linear :func:`walk_commits`, this follows ``parent2_commit_id`` |
| 223 | at merge commits so events on merged feature branches are never missed. |
| 224 | |
| 225 | Args: |
| 226 | root: Repository root. |
| 227 | start_commit_id: SHA-256 of the commit to start from. |
| 228 | max_commits: Safety cap — returns ``truncated=True`` if reached. |
| 229 | stop_at_commit_id: When set, stop BFS *before* entering this commit |
| 230 | (exclusive lower bound, useful for range queries). |
| 231 | |
| 232 | Returns: |
| 233 | ``(commits, truncated)`` — list sorted newest-first by ``committed_at``; |
| 234 | ``truncated=True`` when ``max_commits`` was hit before exhausting the DAG. |
| 235 | """ |
| 236 | commits_by_id: _CommitIndex = {} |
| 237 | queue: deque[str] = deque([start_commit_id]) |
| 238 | seen: set[str] = set() |
| 239 | truncated = False |
| 240 | |
| 241 | while queue: |
| 242 | if len(commits_by_id) >= max_commits: |
| 243 | truncated = True |
| 244 | break |
| 245 | commit_id = queue.popleft() |
| 246 | if commit_id in seen: |
| 247 | continue |
| 248 | if stop_at_commit_id and commit_id == stop_at_commit_id: |
| 249 | continue |
| 250 | seen.add(commit_id) |
| 251 | commit = read_commit(root, commit_id) |
| 252 | if commit is None: |
| 253 | continue |
| 254 | commits_by_id[commit_id] = commit |
| 255 | if commit.parent_commit_id: |
| 256 | queue.append(commit.parent_commit_id) |
| 257 | if commit.parent2_commit_id: |
| 258 | queue.append(commit.parent2_commit_id) |
| 259 | |
| 260 | return ( |
| 261 | sorted(commits_by_id.values(), key=lambda c: c.committed_at, reverse=True), |
| 262 | truncated, |
| 263 | ) |
| 264 | |
| 265 | |
| 266 | # --------------------------------------------------------------------------- |
| 267 | # Op traversal helpers |
| 268 | # --------------------------------------------------------------------------- |
| 269 | |
| 270 | |
| 271 | def flat_symbol_ops(ops: list[DomainOp]) -> Iterator[DomainOp]: |
| 272 | """Yield all leaf ops, recursing into PatchOp.child_ops. |
| 273 | |
| 274 | Only yields ops that have a symbol-level address (i.e. contain ``::``). |
| 275 | """ |
| 276 | for op in ops: |
| 277 | if op["op"] == "patch": |
| 278 | for child in op["child_ops"]: |
| 279 | if "::" in child["address"]: |
| 280 | yield child |
| 281 | elif "::" in op["address"]: |
| 282 | yield op |
| 283 | |
| 284 | |
| 285 | def touched_files(ops: list[DomainOp]) -> frozenset[str]: |
| 286 | """Return the set of file paths that appear as PatchOp addresses in *ops*. |
| 287 | |
| 288 | Only counts files that had symbol-level child ops (semantic changes), |
| 289 | not coarse file-level replace/insert/delete ops. |
| 290 | """ |
| 291 | files: set[str] = set() |
| 292 | for op in ops: |
| 293 | if op["op"] == "patch" and op["child_ops"]: |
| 294 | files.add(op["address"]) |
| 295 | return frozenset(files) |
| 296 | |
| 297 | |
| 298 | def file_pairs(files: frozenset[str]) -> Iterator[tuple[str, str]]: |
| 299 | """Yield all ordered pairs ``(a, b)`` with ``a < b`` from *files*.""" |
| 300 | yield from itertools.combinations(sorted(files), 2) |