gabriel / muse public
query.py python
557 lines 18.4 KB
Raw
sha256:660fcac1df3ab28f61862e961890bd2ca8b754fa0242079d93ca1e25037ec8a6 chore(tests): add docstring to tests/__init__.py so rc14 tr… Human 26 days ago
1 """muse code query — symbol graph predicate query (v2).
2
3 SQL for your codebase. A full predicate DSL over the typed, content-addressed
4 symbol graph — with OR, NOT, grouping, and an expanded field set.
5
6 v2 grammar::
7
8 expr = or_expr
9 or_expr = and_expr ( OR and_expr )*
10 and_expr = not_expr ( [AND] not_expr )* # implicit AND
11 not_expr = NOT primary | primary
12 primary = "(" expr ")" | atom
13 atom = KEY OP VALUE
14
15 Supported operators::
16
17 = exact match
18 ~= contains (case-insensitive)
19 ^= starts with (case-insensitive)
20 $= ends with (case-insensitive)
21 != not equal
22
23 Supported keys::
24
25 kind function | class | method | variable | import | …
26 language Python | Go | Rust | TypeScript | …
27 name bare symbol name
28 qualified_name dotted name (User.save)
29 file file path
30 hash content_id prefix (exact-body match)
31 body_hash body_hash prefix
32 signature_id signature_id prefix
33 lineno_gt symbol starts after line N
34 lineno_lt symbol starts before line N
35 size_gt symbol body exceeds N lines (end_lineno − lineno > N)
36 size_lt symbol body shorter than N lines
37
38 Usage::
39
40 muse code query "kind=function" "language=Python" "name~=validate"
41 muse code query "(kind=function OR kind=method) name^=_"
42 muse code query "NOT kind=import" "file~=billing"
43 muse code query "hash=a3f2c9"
44 muse code query "kind=function" "name$=_test" --commit HEAD~10
45 muse code query "kind=function" "name~=validate" --all-commits
46 muse code query "kind=function" "size_gt=50" --sort size # biggest fns
47 muse code query "kind=function" "name~=compute" --count # just the count
48 muse code query "kind=function" --unique-bodies # find clones
49 muse code query "kind=function" --all-commits --since 2026-01-01 # added this year
50 """
51
52 from __future__ import annotations
53
54 import argparse
55 import collections.abc
56 import datetime
57 import json
58 import logging
59 import pathlib
60 import sys
61
62 from muse._version import __version__
63 from muse.core.errors import ExitCode
64 from muse.core.repo import parse_date_arg, read_repo_id, require_repo
65 from muse.core.store import (
66 CommitRecord,
67 get_all_commits,
68 get_commit_snapshot_manifest,
69 read_current_branch,
70 resolve_commit_ref,
71 )
72 from muse.core.symbol_cache import SymbolCache, load_symbol_cache
73 from muse.plugins.code._predicate import Predicate, PredicateError, parse_query
74 from muse.plugins.code._query import language_of, symbols_for_snapshot
75 from muse.plugins.code.ast_parser import SymbolRecord
76 from muse.core.validation import clamp_int, sanitize_display
77
78
79
80
81 type _QueryResult = dict[str, str | int | bool]
82 type _StrMap = dict[str, str]
83 type _IconMap = dict[str, str]
84 logger = logging.getLogger(__name__)
85
86 _KIND_ICON: _IconMap = {
87 "function": "fn",
88 "async_function": "fn~",
89 "class": "class",
90 "method": "method",
91 "async_method": "method~",
92 "variable": "var",
93 "import": "import",
94 }
95
96 _VALID_SORT_FIELDS = frozenset({"file", "name", "kind", "lineno", "size"})
97
98
99
100 class _HistoricalMatch:
101 """A symbol match found in a historical commit (--all-commits mode)."""
102
103 def __init__(
104 self,
105 address: str,
106 rec: SymbolRecord,
107 commit: CommitRecord,
108 first_seen: bool,
109 ) -> None:
110 self.address = address
111 self.rec = rec
112 self.commit = commit
113 self.first_seen = first_seen
114
115 def to_dict(self) -> _QueryResult:
116 return {
117 "address": self.address,
118 "kind": self.rec["kind"],
119 "name": self.rec["name"],
120 "content_id": self.rec["content_id"],
121 "first_seen": self.first_seen,
122 "commit_id": self.commit.commit_id,
123 "commit_message": self.commit.message,
124 "committed_at": self.commit.committed_at.isoformat(),
125 "branch": self.commit.branch,
126 }
127
128
129 def _query_all_commits(
130 root: pathlib.Path,
131 filters: list[Predicate],
132 max_commits: int,
133 since: datetime.date | None,
134 until: datetime.date | None,
135 ) -> tuple[list[_HistoricalMatch], bool]:
136 """Walk every commit oldest-first, apply predicates against each snapshot.
137
138 Shares one ``SymbolCache`` instance across all snapshot loads so the cache
139 is read from disk exactly once and written back at most once — instead of
140 once per snapshot. On a warm cache this reduces wall time from O(n×200ms)
141 to O(1×load + n×dict_lookup).
142
143 Deduplicates on ``snapshot_id`` — commits sharing a snapshot (e.g. merge
144 commits with no file changes) are processed exactly once.
145
146 Returns:
147 ``(matches, truncated)`` — ``truncated`` is True when the walk was
148 capped at ``max_commits``.
149 """
150 all_commits = get_all_commits(root)
151 if not all_commits:
152 return [], False
153
154 sorted_commits = sorted(all_commits, key=lambda c: c.committed_at)
155
156 # Apply date filters early to avoid unnecessary snapshot loading.
157 if since is not None:
158 sorted_commits = [
159 c for c in sorted_commits if c.committed_at.date() >= since
160 ]
161 if until is not None:
162 sorted_commits = [
163 c for c in sorted_commits if c.committed_at.date() <= until
164 ]
165
166 truncated = len(sorted_commits) > max_commits
167 sorted_commits = sorted_commits[:max_commits]
168
169 results: list[_HistoricalMatch] = []
170 first_seen_map: _StrMap = {}
171 seen_snapshots: set[str] = set()
172
173 # Load the symbol cache once; share it across all snapshot iterations.
174 shared_cache: SymbolCache = load_symbol_cache(root)
175
176 try:
177 for commit in sorted_commits:
178 # Skip commits whose snapshot was already processed.
179 if commit.snapshot_id in seen_snapshots:
180 continue
181 seen_snapshots.add(commit.snapshot_id)
182
183 manifest = get_commit_snapshot_manifest(root, commit.commit_id) or {}
184 if not manifest:
185 continue
186
187 symbol_map = symbols_for_snapshot(root, manifest, cache=shared_cache)
188 for file_path, tree in sorted(symbol_map.items()):
189 for addr, rec in sorted(tree.items(), key=lambda kv: kv[1]["lineno"]):
190 if not all(f(file_path, rec) for f in filters):
191 continue
192 cid = rec["content_id"]
193 is_first = cid not in first_seen_map
194 if is_first:
195 first_seen_map[cid] = commit.commit_id
196 results.append(_HistoricalMatch(addr, rec, commit, is_first))
197 finally:
198 # Persist any newly parsed entries even if we exit early.
199 shared_cache.save()
200
201 return results, truncated
202
203
204 _SortTuple = tuple[str, str, SymbolRecord]
205
206
207 def _sort_key(sort_by: str) -> collections.abc.Callable[[_SortTuple], tuple[str | int, ...]]:
208 """Return a sort key function for a list of ``(file_path, addr, rec)`` tuples."""
209 if sort_by == "name":
210 return lambda t: (t[2]["name"].lower(), t[0], t[2]["lineno"])
211 if sort_by == "kind":
212 return lambda t: (t[2]["kind"], t[0], t[2]["lineno"])
213 if sort_by == "lineno":
214 return lambda t: (t[2]["lineno"], t[0])
215 if sort_by == "size":
216 # Negate size so largest comes first.
217 return lambda t: (-(t[2]["end_lineno"] - t[2]["lineno"]), t[0])
218 # Default: file then lineno.
219 return lambda t: (t[0], t[2]["lineno"])
220
221
222 def register(
223 subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]",
224 ) -> None:
225 """Register the query subcommand."""
226 parser = subparsers.add_parser(
227 "query",
228 help="Query the symbol graph with a predicate DSL.",
229 description=__doc__,
230 formatter_class=argparse.RawDescriptionHelpFormatter,
231 )
232 parser.add_argument(
233 "predicates",
234 nargs="+",
235 metavar="PREDICATE",
236 help='One or more predicates, e.g. "kind=function" "name~=validate".',
237 )
238 parser.add_argument(
239 "--commit", "-c",
240 dest="ref",
241 default=None,
242 metavar="REF",
243 help="Query a historical snapshot instead of HEAD.",
244 )
245 parser.add_argument(
246 "--all-commits",
247 action="store_true",
248 help=(
249 "Search across ALL commits (every branch). Enables temporal"
250 " hash= queries: find when a function body first appeared."
251 " Mutually exclusive with --commit."
252 ),
253 )
254 parser.add_argument(
255 "--since",
256 metavar="YYYY-MM-DD",
257 default=None,
258 help=(
259 "With --all-commits: only consider commits on or after this date."
260 ),
261 )
262 parser.add_argument(
263 "--until",
264 metavar="YYYY-MM-DD",
265 default=None,
266 help=(
267 "With --all-commits: only consider commits on or before this date."
268 ),
269 )
270 parser.add_argument(
271 "--max-commits",
272 type=int,
273 default=10_000,
274 metavar="N",
275 help=(
276 "With --all-commits: cap the number of commits walked"
277 " (default: 10000)."
278 ),
279 )
280 parser.add_argument(
281 "--limit",
282 type=int,
283 default=0,
284 metavar="N",
285 help="Cap the number of results returned (0 = unlimited).",
286 )
287 parser.add_argument(
288 "--sort",
289 default="file",
290 metavar="FIELD",
291 choices=sorted(_VALID_SORT_FIELDS),
292 help=(
293 f"Sort results by field: {', '.join(sorted(_VALID_SORT_FIELDS))}"
294 " (default: file)."
295 ),
296 )
297 parser.add_argument(
298 "--count",
299 action="store_true",
300 help="Print only the count of matching symbols — no symbol list.",
301 )
302 parser.add_argument(
303 "--unique-bodies",
304 action="store_true",
305 help=(
306 "Deduplicate by content_id — show only unique implementations."
307 " Turns muse query into a clone detector."
308 ),
309 )
310 parser.add_argument(
311 "--hashes",
312 dest="show_hashes",
313 action="store_true",
314 help="Include content hashes in output.",
315 )
316 parser.add_argument(
317 "--json",
318 dest="as_json",
319 action="store_true",
320 help="Emit results as JSON.",
321 )
322 parser.set_defaults(func=run)
323
324
325 def run(args: argparse.Namespace) -> None:
326 """Query the symbol graph with a predicate DSL.
327
328 ``muse query`` is SQL for your codebase. Every predicate is evaluated
329 against the typed, content-addressed symbol graph — not raw text.
330
331 New in v2.1:
332 ``size_gt=N`` / ``size_lt=N`` — filter by symbol body line count.
333 ``--count`` — emit only the result count.
334 ``--limit N`` — cap results (like SQL LIMIT).
335 ``--sort FIELD`` — sort by file, name, kind, lineno, or size.
336 ``--unique-bodies`` — deduplicate by content_id (clone detector mode).
337 ``--since / --until YYYY-MM-DD`` — temporal range for --all-commits.
338 """
339 predicates: list[str] = args.predicates
340 ref: str | None = args.ref
341 all_commits: bool = args.all_commits
342 show_hashes: bool = args.show_hashes
343 as_json: bool = args.as_json
344 count_only: bool = args.count
345 limit: int = clamp_int(args.limit, 0, 10000, 'limit')
346 sort_by: str = args.sort
347 unique_bodies: bool = args.unique_bodies
348 max_commits: int = clamp_int(args.max_commits, 1, 100000, 'max_commits')
349
350 root = require_repo()
351 repo_id = read_repo_id(root)
352 branch = read_current_branch(root)
353
354 if not predicates:
355 print("❌ At least one predicate is required.", file=sys.stderr)
356 raise SystemExit(ExitCode.USER_ERROR)
357
358 if all_commits and ref is not None:
359 print(
360 "❌ --all-commits and --commit are mutually exclusive.",
361 file=sys.stderr,
362 )
363 raise SystemExit(ExitCode.USER_ERROR)
364
365 if limit < 0:
366 print("❌ --limit must be >= 0.", file=sys.stderr)
367 raise SystemExit(ExitCode.USER_ERROR)
368
369 if max_commits < 1:
370 print("❌ --max-commits must be >= 1.", file=sys.stderr)
371 raise SystemExit(ExitCode.USER_ERROR)
372
373 # Parse --since / --until date filters.
374 since_date: datetime.date | None = (
375 parse_date_arg(args.since, "--since").date() if args.since else None
376 )
377 until_date: datetime.date | None = (
378 parse_date_arg(args.until, "--until").date() if args.until else None
379 )
380
381 if (since_date or until_date) and not all_commits:
382 print(
383 "❌ --since / --until require --all-commits.", file=sys.stderr
384 )
385 raise SystemExit(ExitCode.USER_ERROR)
386
387 # Parse predicates via the v2 grammar.
388 try:
389 combined_predicate: Predicate = parse_query(predicates)
390 except PredicateError as exc:
391 print(f"❌ {exc}", file=sys.stderr)
392 raise SystemExit(ExitCode.USER_ERROR)
393 filters: list[Predicate] = [combined_predicate]
394
395 # ── --all-commits mode ────────────────────────────────────────────────────
396 if all_commits:
397 historical, truncated = _query_all_commits(
398 root, filters, max_commits, since_date, until_date
399 )
400
401 if as_json:
402 print(
403 json.dumps(
404 {
405 "schema_version": __version__,
406 "mode": "all-commits",
407 "truncated": truncated,
408 "results": [h.to_dict() for h in historical],
409 },
410 indent=2,
411 )
412 )
413 return
414
415 if not historical:
416 pred_display = " AND ".join(predicates)
417 print(
418 f" (no symbols matching: {pred_display}"
419 f" [searched all commits])"
420 )
421 return
422
423 # Deduplicate for display: show unique addresses with first-seen commit.
424 seen_addrs: set[str] = set()
425 unique: list[_HistoricalMatch] = []
426 for h in historical:
427 if h.first_seen and h.address not in seen_addrs:
428 seen_addrs.add(h.address)
429 unique.append(h)
430
431 if limit > 0:
432 unique = unique[:limit]
433
434 if count_only:
435 print(len(unique))
436 return
437
438 pred_display = " AND ".join(predicates)
439 trunc_note = " ⚠️ truncated" if truncated else ""
440 print(
441 f"\n{len(unique)} unique symbol(s) matching"
442 f" [{pred_display}] across all commits{trunc_note}\n"
443 )
444 for h in unique:
445 date_str = h.commit.committed_at.strftime("%Y-%m-%d")
446 short_id = h.commit.commit_id[:8]
447 icon = _KIND_ICON.get(h.rec["kind"], h.rec["kind"])
448 hash_part = f" {h.rec['content_id'][:8]}.." if show_hashes else ""
449 branch_label = (
450 f" [{h.commit.branch}]" if h.commit.branch else ""
451 )
452 print(
453 f" {h.address:<60} {icon:<8}"
454 f" first seen {short_id} {date_str}"
455 f"{branch_label}{hash_part}"
456 )
457 return
458
459 # ── Single-snapshot mode (default) ────────────────────────────────────────
460 commit = resolve_commit_ref(root, repo_id, branch, ref)
461 if commit is None:
462 print(
463 f"❌ Commit '{ref or 'HEAD'}' not found.", file=sys.stderr
464 )
465 raise SystemExit(ExitCode.USER_ERROR)
466
467 manifest = get_commit_snapshot_manifest(root, commit.commit_id) or {}
468 symbol_map = symbols_for_snapshot(root, manifest)
469
470 # Collect matches.
471 matches: list[tuple[str, str, SymbolRecord]] = []
472 for file_path, tree in symbol_map.items():
473 for addr, rec in tree.items():
474 if all(f(file_path, rec) for f in filters):
475 matches.append((file_path, addr, rec))
476
477 # Sort.
478 matches.sort(key=_sort_key(sort_by))
479
480 # Unique-bodies: deduplicate by content_id.
481 if unique_bodies:
482 seen_cids: set[str] = set()
483 deduped: list[tuple[str, str, SymbolRecord]] = []
484 for fp, addr, rec in matches:
485 cid = rec["content_id"]
486 if cid not in seen_cids:
487 seen_cids.add(cid)
488 deduped.append((fp, addr, rec))
489 matches = deduped
490
491 # Apply limit.
492 limited = limit > 0 and len(matches) > limit
493 if limited:
494 matches = matches[:limit]
495
496 # Count-only output.
497 if count_only:
498 print(len(matches))
499 return
500
501 # JSON output.
502 if as_json:
503 out: list[dict[str, str | int]] = []
504 for fp, addr, rec in matches:
505 out.append(
506 {
507 "address": addr,
508 "kind": rec["kind"],
509 "name": rec["name"],
510 "qualified_name": rec["qualified_name"],
511 "file": fp,
512 "lineno": rec["lineno"],
513 "end_lineno": rec["end_lineno"],
514 "size": rec["end_lineno"] - rec["lineno"],
515 "language": language_of(fp),
516 "content_id": rec["content_id"],
517 "body_hash": rec["body_hash"],
518 "signature_id": rec["signature_id"],
519 }
520 )
521 print(
522 json.dumps(
523 {
524 "schema_version": __version__,
525 "commit": commit.commit_id[:8],
526 "sort": sort_by,
527 "unique_bodies": unique_bodies,
528 "truncated": limited,
529 "results": out,
530 },
531 indent=2,
532 )
533 )
534 return
535
536 # Human-readable output.
537 if not matches:
538 pred_str = " AND ".join(predicates)
539 print(f" (no symbols matching: {pred_str})")
540 return
541
542 files_seen: set[str] = set()
543 for fp, addr, rec in matches:
544 files_seen.add(fp)
545 icon = _KIND_ICON.get(rec["kind"], rec["kind"])
546 line = rec["lineno"]
547 size = rec["end_lineno"] - rec["lineno"]
548 hash_part = f" {rec['content_id'][:8]}.." if show_hashes else ""
549 size_part = f" {size:>3}L" if sort_by == "size" else ""
550 print(f" {sanitize_display(addr):<60} {icon:<10} line {line:>4}{size_part}{hash_part}")
551
552 pred_display = " AND ".join(predicates)
553 trunc_note = f" (limited to {limit})" if limited else ""
554 print(
555 f"\n{len(matches)} match(es) across {len(files_seen)} file(s)"
556 f" [{pred_display}]{trunc_note}"
557 )
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