_query.py
python
sha256:24d9cb81ac1dd071e5d98862c740ad101180588361b4f96a790af79a99aed901
docs: architectural plan for unified-store snapshot-content…
Sonnet 5
10 days ago
| 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 | import itertools |
| 13 | import logging |
| 14 | import pathlib |
| 15 | import re |
| 16 | from collections.abc import Iterator |
| 17 | |
| 18 | from muse.core.types import Manifest, Metadata, blob_id, short_id |
| 19 | from muse.core.object_store import read_object |
| 20 | from muse.core.stat_cache import StatCache |
| 21 | from muse.core.commits import ( |
| 22 | CommitRecord, |
| 23 | read_commit, |
| 24 | ) |
| 25 | from muse.core.symbol_cache import SymbolCache, load_symbol_cache |
| 26 | from muse.domain import DomainOp |
| 27 | from muse.plugins.code.ast_parser import ( |
| 28 | SymbolTree, |
| 29 | parse_symbols, |
| 30 | ) |
| 31 | |
| 32 | type _SymbolTreeMap = dict[str, SymbolTree] |
| 33 | type _CommitIndex = dict[str, CommitRecord] |
| 34 | |
| 35 | logger = logging.getLogger(__name__) |
| 36 | |
| 37 | # --------------------------------------------------------------------------- |
| 38 | # Language classification |
| 39 | # --------------------------------------------------------------------------- |
| 40 | |
| 41 | _SUFFIX_LANG: Metadata = { |
| 42 | ".py": "Python", ".pyi": "Python", |
| 43 | ".ts": "TypeScript", ".tsx": "TypeScript", |
| 44 | ".js": "JavaScript", ".jsx": "JavaScript", |
| 45 | ".mjs": "JavaScript", ".cjs": "JavaScript", |
| 46 | ".go": "Go", |
| 47 | ".rs": "Rust", |
| 48 | ".java": "Java", |
| 49 | ".cs": "C#", |
| 50 | ".c": "C", ".h": "C", |
| 51 | ".cpp": "C++", ".cc": "C++", ".cxx": "C++", ".hpp": "C++", ".hxx": "C++", |
| 52 | ".rb": "Ruby", |
| 53 | ".kt": "Kotlin", ".kts": "Kotlin", |
| 54 | ".swift": "Swift", |
| 55 | ".md": "Markdown", ".rst": "reStructuredText", ".txt": "Text", |
| 56 | ".toml": "TOML", |
| 57 | ".yaml": "YAML", ".yml": "YAML", |
| 58 | ".json": "JSON", ".jsonc": "JSON", |
| 59 | ".css": "CSS", ".scss": "SCSS", |
| 60 | ".html": "HTML", ".htm": "HTML", |
| 61 | ".sql": "SQL", |
| 62 | ".sh": "Shell", ".bash": "Shell", ".zsh": "Shell", |
| 63 | ".proto": "Protobuf", |
| 64 | ".tf": "Terraform", |
| 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 | def is_semantic(file_path: str) -> bool: |
| 73 | """Return ``True`` if *file_path* has a suffix with AST-level support.""" |
| 74 | from muse.plugins.code.ast_parser import SEMANTIC_EXTENSIONS |
| 75 | suffix = pathlib.PurePosixPath(file_path).suffix.lower() |
| 76 | return suffix in SEMANTIC_EXTENSIONS |
| 77 | |
| 78 | # --------------------------------------------------------------------------- |
| 79 | # Language normalisation |
| 80 | # --------------------------------------------------------------------------- |
| 81 | |
| 82 | _LANG_CANONICAL: dict[str, str] = {lang.lower(): lang for lang in set(_SUFFIX_LANG.values())} |
| 83 | |
| 84 | def normalise_language(lang: str) -> str: |
| 85 | """Normalise a user-supplied language name to its canonical display form. |
| 86 | |
| 87 | Case-insensitive. Returns the canonical capitalisation (e.g. ``"python"`` |
| 88 | → ``"Python"``) or the stripped input unchanged if unrecognised. |
| 89 | """ |
| 90 | return _LANG_CANONICAL.get(lang.strip().lower(), lang.strip()) |
| 91 | |
| 92 | # --------------------------------------------------------------------------- |
| 93 | # Test-file detection |
| 94 | # --------------------------------------------------------------------------- |
| 95 | |
| 96 | _TEST_PATTERNS: tuple[re.Pattern[str], ...] = ( |
| 97 | re.compile(r"(^|/)test_"), |
| 98 | re.compile(r"_test\.py$"), |
| 99 | re.compile(r"(^|/)tests/"), |
| 100 | re.compile(r"(^|/)spec/"), |
| 101 | re.compile(r"(^|/)conftest\.py$"), |
| 102 | ) |
| 103 | |
| 104 | def is_test_file(file_path: str) -> bool: |
| 105 | """Return ``True`` if *file_path* looks like a test or conftest file.""" |
| 106 | return any(p.search(file_path) for p in _TEST_PATTERNS) |
| 107 | |
| 108 | # --------------------------------------------------------------------------- |
| 109 | # Symbol extraction from a snapshot manifest |
| 110 | # --------------------------------------------------------------------------- |
| 111 | |
| 112 | def _rekey_tree(tree: SymbolTree, file_path: str) -> SymbolTree: |
| 113 | """Re-key *tree* so every address starts with *file_path*. |
| 114 | |
| 115 | The ``SymbolCache`` is keyed by content SHA-256. When two files have |
| 116 | identical bytes they share one cache entry whose addresses carry the path |
| 117 | of whichever file was parsed first. This function patches those addresses |
| 118 | to the actual *file_path* being processed — an O(n) pass that is only |
| 119 | triggered on a cache hit where the prefix differs. |
| 120 | """ |
| 121 | first = next(iter(tree), None) |
| 122 | if first is None: |
| 123 | return tree |
| 124 | cached_prefix = first.split("::")[0] if "::" in first else None |
| 125 | if cached_prefix == file_path: |
| 126 | return tree # Addresses already correct — common case. |
| 127 | return { |
| 128 | f"{file_path}::{addr.split('::', 1)[1]}" if "::" in addr else addr: rec |
| 129 | for addr, rec in tree.items() |
| 130 | } |
| 131 | |
| 132 | def symbols_for_snapshot( |
| 133 | root: pathlib.Path, |
| 134 | manifest: Manifest, |
| 135 | *, |
| 136 | kind_filter: str | None = None, |
| 137 | file_filter: str | None = None, |
| 138 | language_filter: str | None = None, |
| 139 | workdir: pathlib.Path | None = None, |
| 140 | cache: SymbolCache | None = None, |
| 141 | stat_cache: StatCache | None = None, |
| 142 | ) -> _SymbolTreeMap: |
| 143 | """Extract symbol trees for all semantic files in *manifest*. |
| 144 | |
| 145 | Results are served from the persistent symbol cache when available, |
| 146 | cutting a full-snapshot scan from ~1,300 ms to ~22 ms on a warm cache. |
| 147 | The cache is keyed by the SHA-256 of the file bytes (``object_id`` for |
| 148 | committed files; freshly computed SHA-256 for working-tree reads), so |
| 149 | every hit is guaranteed correct. |
| 150 | |
| 151 | Args: |
| 152 | root: Repository root (used to locate the object store and |
| 153 | load/save the symbol cache). |
| 154 | manifest: Snapshot manifest mapping file path → SHA-256. |
| 155 | kind_filter: If set, only include symbols with this ``kind``. |
| 156 | file_filter: If set, only include symbols from this exact file path. |
| 157 | language_filter: If set, only include symbols from files of this language. |
| 158 | workdir: When set, file content is read from *workdir* first |
| 159 | (working-tree mode), falling back to the object store if |
| 160 | the file does not exist on disk. Pass ``root`` to get |
| 161 | live, uncommitted changes reflected immediately. |
| 162 | cache: Optional pre-loaded ``SymbolCache`` instance. When |
| 163 | ``None`` the cache is loaded from disk automatically |
| 164 | and saved at the end of the call. Pass an explicit |
| 165 | instance when batching multiple calls to share one |
| 166 | cache load/save cycle. |
| 167 | stat_cache: Optional ``StatCache`` instance. When supplied and |
| 168 | ``workdir`` is set, the SHA-256 cache key for each |
| 169 | disk file is derived from ``(ino, mtime, size)`` |
| 170 | without reading the file bytes. On a warm |
| 171 | ``SymbolCache`` hit this means zero file reads for |
| 172 | unchanged files — the dominant cost of the working-tree |
| 173 | diff path. Has no effect when ``workdir`` is ``None``. |
| 174 | |
| 175 | Returns: |
| 176 | Dict mapping ``file_path → SymbolTree``; empty trees are omitted. |
| 177 | """ |
| 178 | if cache is None: |
| 179 | active_cache: SymbolCache = load_symbol_cache(root) |
| 180 | own_cache = True |
| 181 | else: |
| 182 | active_cache = cache |
| 183 | own_cache = False |
| 184 | |
| 185 | result: _SymbolTreeMap = {} |
| 186 | for file_path, object_id in sorted(manifest.items()): |
| 187 | if not is_semantic(file_path): |
| 188 | continue |
| 189 | if file_filter and file_path != file_filter: |
| 190 | continue |
| 191 | if language_filter and language_of(file_path) != language_filter: |
| 192 | continue |
| 193 | |
| 194 | # --- Determine cache key and read bytes if needed ---------------- |
| 195 | # For working-tree files the object_id is the committed version; |
| 196 | # the disk content may differ. We compute the SHA-256 of whatever |
| 197 | # bytes we actually parse so the cache key is always content-addressed. |
| 198 | cache_key = object_id |
| 199 | raw: bytes | None = None |
| 200 | # When stat_cache is active we defer reading disk bytes until we know |
| 201 | # whether the SymbolCache will miss. This variable holds the path to |
| 202 | # read from if a sym-cache miss forces us to load and parse the file. |
| 203 | _disk_for_miss: pathlib.Path | None = None |
| 204 | |
| 205 | if workdir is not None: |
| 206 | disk_path = workdir / file_path |
| 207 | if disk_path.is_file(): |
| 208 | try: |
| 209 | if stat_cache is not None: |
| 210 | st = disk_path.stat() |
| 211 | # get_cached returns the SHA-256 (sha256:… prefix). |
| 212 | # On a stat-cache hit it does so without reading bytes. |
| 213 | # On a stat-cache miss it reads and hashes internally. |
| 214 | cache_key = stat_cache.get_cached( |
| 215 | file_path, str(disk_path), |
| 216 | st.st_mtime, st.st_size, st.st_ino, |
| 217 | ) |
| 218 | _disk_for_miss = disk_path # bytes deferred |
| 219 | else: |
| 220 | raw = disk_path.read_bytes() |
| 221 | # Recompute the key for the actual disk content. |
| 222 | cache_key = blob_id(raw) |
| 223 | except OSError as exc: |
| 224 | logger.debug("Could not read %s from workdir: %s", file_path, exc) |
| 225 | |
| 226 | # --- Cache lookup ------------------------------------------------ |
| 227 | tree = active_cache.get(cache_key) |
| 228 | if tree is None: |
| 229 | # Cache miss — fetch bytes and parse. |
| 230 | if raw is None: |
| 231 | if _disk_for_miss is not None: |
| 232 | # stat_cache path: bytes were not pre-read; load them now. |
| 233 | try: |
| 234 | raw = _disk_for_miss.read_bytes() |
| 235 | except OSError as exc: |
| 236 | logger.debug("Could not read %s from disk: %s", file_path, exc) |
| 237 | if raw is None: |
| 238 | raw = read_object(root, object_id) |
| 239 | if raw is None: |
| 240 | logger.debug("Object %s missing for %s — skipping", short_id(object_id), file_path) |
| 241 | continue |
| 242 | tree = parse_symbols(raw, file_path) |
| 243 | active_cache.put(cache_key, tree) |
| 244 | else: |
| 245 | # Cache hit — re-key addresses when two files share the same |
| 246 | # SHA-256 (identical content). The cache stores the tree built |
| 247 | # for whichever file was parsed first; its addresses carry that |
| 248 | # file's path prefix. Patch them to the current file_path so |
| 249 | # cross-file clone detection and all downstream consumers get |
| 250 | # correct symbol addresses. |
| 251 | tree = _rekey_tree(tree, file_path) |
| 252 | |
| 253 | # --- Apply filters on top of the (potentially cached) full tree -- |
| 254 | if kind_filter: |
| 255 | tree = {addr: rec for addr, rec in tree.items() if rec["kind"] == kind_filter} |
| 256 | if tree: |
| 257 | result[file_path] = tree |
| 258 | |
| 259 | if own_cache: |
| 260 | active_cache.save() |
| 261 | |
| 262 | return result |
| 263 | |
| 264 | # --------------------------------------------------------------------------- |
| 265 | # Commit-graph walking |
| 266 | # --------------------------------------------------------------------------- |
| 267 | |
| 268 | def walk_commits_bfs( |
| 269 | root: pathlib.Path, |
| 270 | start_commit_id: str, |
| 271 | max_commits: int = 500, |
| 272 | stop_at_commit_id: str | None = None, |
| 273 | ) -> tuple[list[CommitRecord], bool]: |
| 274 | """BFS walk of the commit DAG from *start_commit_id*, newest-first. |
| 275 | |
| 276 | Unlike the linear :func:`walk_commits`, this follows ``parent2_commit_id`` |
| 277 | at merge commits so events on merged feature branches are never missed. |
| 278 | |
| 279 | Args: |
| 280 | root: Repository root. |
| 281 | start_commit_id: SHA-256 of the commit to start from. |
| 282 | max_commits: Safety cap — returns ``truncated=True`` if reached. |
| 283 | stop_at_commit_id: When set, stop BFS *before* entering this commit |
| 284 | (exclusive lower bound, useful for range queries). |
| 285 | |
| 286 | Returns: |
| 287 | ``(commits, truncated)`` — list sorted newest-first by ``committed_at``; |
| 288 | ``truncated=True`` when ``max_commits`` was hit before exhausting the DAG. |
| 289 | """ |
| 290 | from muse.core.graph import iter_ancestors |
| 291 | |
| 292 | prune = (lambda cid: cid == stop_at_commit_id) if stop_at_commit_id else None |
| 293 | truncated = False |
| 294 | commits_by_id: _CommitIndex = {} |
| 295 | |
| 296 | for commit in iter_ancestors(root, start_commit_id, prune=prune): |
| 297 | if len(commits_by_id) >= max_commits: |
| 298 | truncated = True |
| 299 | break |
| 300 | commits_by_id[commit.commit_id] = commit |
| 301 | |
| 302 | return ( |
| 303 | sorted(commits_by_id.values(), key=lambda c: c.committed_at, reverse=True), |
| 304 | truncated, |
| 305 | ) |
| 306 | |
| 307 | # --------------------------------------------------------------------------- |
| 308 | # Op traversal helpers |
| 309 | # --------------------------------------------------------------------------- |
| 310 | |
| 311 | def dir_of(path: str) -> str: |
| 312 | """Return the immediate parent directory of *path* as a POSIX string. |
| 313 | |
| 314 | - ``"src/billing.py"`` → ``"src"`` |
| 315 | - ``"muse/cli/commands/cat.py"`` → ``"muse/cli/commands"`` |
| 316 | - ``"main.py"`` → ``"."`` |
| 317 | - ``"src/"`` → ``"src"`` (trailing slash stripped first) |
| 318 | """ |
| 319 | if path.endswith("/"): |
| 320 | return path.rstrip("/") or "." |
| 321 | parent = str(pathlib.PurePosixPath(path).parent) |
| 322 | return parent |
| 323 | |
| 324 | |
| 325 | def flat_directory_ops(ops: list[DomainOp]) -> Iterator[DomainOp]: |
| 326 | """Yield directory-level ops from *ops*. |
| 327 | |
| 328 | Yields: |
| 329 | - ``RenameOp`` (op == "rename") |
| 330 | - ``InsertOp`` / ``DeleteOp`` whose address ends with "/" or whose |
| 331 | content_summary starts with "directory:" |
| 332 | |
| 333 | Skips file-level and symbol-level ops. |
| 334 | """ |
| 335 | for op in ops: |
| 336 | op_type = op["op"] |
| 337 | if op_type == "rename": |
| 338 | yield op |
| 339 | elif op_type in ("insert", "delete"): |
| 340 | addr: str = op["address"] |
| 341 | summary: str = op.get("content_summary", "") # type: ignore[arg-type] |
| 342 | if addr.endswith("/") or summary.startswith("directory:"): |
| 343 | yield op |
| 344 | |
| 345 | |
| 346 | def touched_directories(ops: list[DomainOp]) -> frozenset[str]: |
| 347 | """Return directories affected by the ops in *ops*. |
| 348 | |
| 349 | Counts: |
| 350 | - The parent dir of every file with semantic child ops in a PatchOp. |
| 351 | - Both ``address`` and ``from_address`` of RenameOps. |
| 352 | |
| 353 | Ignores PatchOps with empty child_ops (no semantic change). |
| 354 | """ |
| 355 | dirs: set[str] = set() |
| 356 | for op in ops: |
| 357 | op_type = op["op"] |
| 358 | if op_type == "patch": |
| 359 | if op["child_ops"]: # type: ignore[index] |
| 360 | dirs.add(dir_of(op["address"])) # type: ignore[arg-type] |
| 361 | elif op_type == "rename": |
| 362 | dirs.add(op["address"]) # type: ignore[arg-type] |
| 363 | dirs.add(op["from_address"]) # type: ignore[union-attr] |
| 364 | return frozenset(dirs) |
| 365 | |
| 366 | |
| 367 | def flat_symbol_ops(ops: list[DomainOp]) -> Iterator[DomainOp]: |
| 368 | """Yield all leaf ops, recursing into PatchOp.child_ops. |
| 369 | |
| 370 | Only yields ops that have a symbol-level address (i.e. contain ``::``). |
| 371 | """ |
| 372 | for op in ops: |
| 373 | if op["op"] == "patch": |
| 374 | for child in op["child_ops"]: |
| 375 | if "::" in child["address"]: |
| 376 | yield child |
| 377 | elif "::" in op["address"]: |
| 378 | yield op |
| 379 | |
| 380 | def touched_files(ops: list[DomainOp]) -> frozenset[str]: |
| 381 | """Return the set of file paths that appear as PatchOp addresses in *ops*. |
| 382 | |
| 383 | Only counts files that had symbol-level child ops (semantic changes), |
| 384 | not coarse file-level replace/insert/delete ops. |
| 385 | """ |
| 386 | files: set[str] = set() |
| 387 | for op in ops: |
| 388 | if op["op"] == "patch" and op["child_ops"]: |
| 389 | files.add(op["address"]) |
| 390 | return frozenset(files) |
| 391 | |
| 392 | def file_pairs(files: frozenset[str]) -> Iterator[tuple[str, str]]: |
| 393 | """Yield all ordered pairs ``(a, b)`` with ``a < b`` from *files*.""" |
| 394 | yield from itertools.combinations(sorted(files), 2) |
File History
1 commit
sha256:24d9cb81ac1dd071e5d98862c740ad101180588361b4f96a790af79a99aed901
docs: architectural plan for unified-store snapshot-content…
Sonnet 5
10 days ago