checkout_symbol.py
file-level
1
files
1
commits
0
hotspots
0
🧊 dead
0
💥 blast risk
| 1 | """muse code checkout-symbol — restore a historical version of a specific symbol. |
| 2 | |
| 3 | Extracts a single named symbol from a historical committed snapshot and writes |
| 4 | it back into the current working-tree file, replacing the current version of |
| 5 | that symbol. |
| 6 | |
| 7 | This is a **surgical** operation: only the target symbol's lines change. |
| 8 | All surrounding code — other symbols, comments, imports, blank lines outside |
| 9 | the symbol boundary — is left untouched. |
| 10 | |
| 11 | Why this matters |
| 12 | ---------------- |
| 13 | Git's ``checkout`` restores entire files. If you need to roll back a single |
| 14 | function while keeping everything else current, you need to manually cherry- |
| 15 | pick lines. ``muse code checkout-symbol`` does this atomically against Muse's |
| 16 | content-addressed symbol index. |
| 17 | |
| 18 | No-op detection: if the symbol body at the target commit is identical to the |
| 19 | current working-tree version (same body hash), the file is not written and |
| 20 | ``"changed": false`` is reported. |
| 21 | |
| 22 | Security note: the file path component of ADDRESS is validated via |
| 23 | ``contain_path()`` before any disk access. Paths that escape the repo root |
| 24 | (e.g. ``../../etc/passwd::foo``) are rejected with exit 1. |
| 25 | |
| 26 | Usage:: |
| 27 | |
| 28 | muse code checkout-symbol "src/billing.py::compute_invoice_total" --commit HEAD~3 |
| 29 | muse code checkout-symbol "src/auth.py::validate_token" --commit abc12345 --dry-run |
| 30 | muse code checkout-symbol "src/billing.py::compute_invoice_total" --commit HEAD~3 --json |
| 31 | |
| 32 | Output (without --dry-run):: |
| 33 | |
| 34 | Restoring: src/billing.py::compute_invoice_total |
| 35 | from commit: abc12345 (2026-02-15) |
| 36 | lines 42–67 → replaced with 31 historical line(s) |
| 37 | ✅ Written to src/billing.py |
| 38 | |
| 39 | Output (no-op — symbol already matches):: |
| 40 | |
| 41 | ✅ src/billing.py::compute_invoice_total already matches commit abc12345 — nothing to do. |
| 42 | |
| 43 | Output (with --dry-run):: |
| 44 | |
| 45 | Dry run — no files will be written. |
| 46 | |
| 47 | Restoring: src/billing.py::compute_invoice_total |
| 48 | from commit: abc12345 (2026-02-15) |
| 49 | |
| 50 | --- current |
| 51 | +++ historical |
| 52 | @@ -42,26 +42,20 @@ |
| 53 | def compute_invoice_total(...): |
| 54 | - ...current body... |
| 55 | + ...historical body... |
| 56 | |
| 57 | JSON output (``--json``):: |
| 58 | |
| 59 | { |
| 60 | "schema_version": "0.1.5", |
| 61 | "address": "src/billing.py::compute_invoice_total", |
| 62 | "file": "src/billing.py", |
| 63 | "branch": "main", |
| 64 | "restored_from": "abc12345", |
| 65 | "dry_run": false, |
| 66 | "changed": true, |
| 67 | "appended": false, |
| 68 | "current_start": 42, |
| 69 | "current_end": 67, |
| 70 | "historical_line_count": 31, |
| 71 | "diff_lines": [] |
| 72 | } |
| 73 | |
| 74 | Flags: |
| 75 | |
| 76 | ``--commit, -c REF`` |
| 77 | Required. Commit to restore from. |
| 78 | |
| 79 | ``--dry-run`` |
| 80 | Print the diff without writing anything. |
| 81 | |
| 82 | ``--json`` |
| 83 | Emit result as JSON for agent consumption. In dry-run mode the JSON |
| 84 | includes a ``diff_lines`` list so agents can inspect what would change. |
| 85 | """ |
| 86 | |
| 87 | from __future__ import annotations |
| 88 | |
| 89 | import argparse |
| 90 | import difflib |
| 91 | import json |
| 92 | import logging |
| 93 | import pathlib |
| 94 | |
| 95 | from muse._version import __version__ |
| 96 | from muse.core.errors import ExitCode |
| 97 | from muse.core.object_store import read_object |
| 98 | from muse.core.repo import read_repo_id, require_repo |
| 99 | from muse.core.store import get_commit_snapshot_manifest, read_current_branch, resolve_commit_ref |
| 100 | from muse.core.validation import contain_path, sanitize_display |
| 101 | from muse.plugins.code.ast_parser import SymbolRecord, SymbolTree, parse_symbols |
| 102 | |
| 103 | logger = logging.getLogger(__name__) |
| 104 | |
| 105 | |
| 106 | # --------------------------------------------------------------------------- |
| 107 | # Internal helpers |
| 108 | # --------------------------------------------------------------------------- |
| 109 | |
| 110 | |
| 111 | def _extract_lines(source: bytes, lineno: int, end_lineno: int) -> list[str]: |
| 112 | """Extract lines *lineno*..*end_lineno* (1-indexed, inclusive). |
| 113 | |
| 114 | Returns an empty list and logs a warning if the range is out of bounds |
| 115 | rather than silently truncating. |
| 116 | """ |
| 117 | all_lines = source.decode("utf-8", errors="replace").splitlines(keepends=True) |
| 118 | total = len(all_lines) |
| 119 | if lineno < 1 or end_lineno < lineno or end_lineno > total: |
| 120 | logger.warning( |
| 121 | "Line range %d–%d out of bounds for %d-line source — returning empty", |
| 122 | lineno, end_lineno, total, |
| 123 | ) |
| 124 | return [] |
| 125 | return all_lines[lineno - 1:end_lineno] |
| 126 | |
| 127 | |
| 128 | def _find_symbol_in_source( |
| 129 | source: bytes, |
| 130 | file_rel: str, |
| 131 | address: str, |
| 132 | ) -> SymbolRecord | None: |
| 133 | """Parse *source* as *file_rel* and return the record for *address*, or None.""" |
| 134 | # file_rel must be the repo-relative path so that parse_symbols builds |
| 135 | # address keys that match the caller's address format (e.g. "src/a.py::fn"). |
| 136 | tree: SymbolTree = parse_symbols(source, file_rel) |
| 137 | return tree.get(address) |
| 138 | |
| 139 | |
| 140 | # --------------------------------------------------------------------------- |
| 141 | # CLI registration and entry point |
| 142 | # --------------------------------------------------------------------------- |
| 143 | |
| 144 | |
| 145 | def register(subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]") -> None: |
| 146 | """Register the checkout-symbol subcommand.""" |
| 147 | parser = subparsers.add_parser( |
| 148 | "checkout-symbol", |
| 149 | help="Restore a historical version of a specific symbol into the working tree.", |
| 150 | description=__doc__, |
| 151 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 152 | ) |
| 153 | parser.add_argument( |
| 154 | "address", metavar="ADDRESS", |
| 155 | help='Symbol address, e.g. "src/billing.py::compute_invoice_total".', |
| 156 | ) |
| 157 | parser.add_argument( |
| 158 | "--commit", "-c", required=True, metavar="REF", dest="ref", |
| 159 | help="Commit to restore the symbol from (required).", |
| 160 | ) |
| 161 | parser.add_argument( |
| 162 | "--dry-run", action="store_true", dest="dry_run", |
| 163 | help="Print the diff without writing anything.", |
| 164 | ) |
| 165 | parser.add_argument( |
| 166 | "--json", action="store_true", dest="as_json", |
| 167 | help="Emit result as JSON for agent consumption.", |
| 168 | ) |
| 169 | parser.set_defaults(func=run) |
| 170 | |
| 171 | |
| 172 | def run(args: argparse.Namespace) -> None: |
| 173 | """Restore a historical version of a specific symbol into the working tree. |
| 174 | |
| 175 | Extracts the symbol body from the given historical commit and splices it |
| 176 | into the current working-tree file at the symbol's current location. |
| 177 | Only the target symbol's lines change; everything else is left untouched. |
| 178 | |
| 179 | If the symbol body already matches the historical version, no file is |
| 180 | written and ``changed=false`` is reported — a safe no-op. |
| 181 | |
| 182 | If the symbol does not exist at ``--commit``, the command exits with an |
| 183 | error. If the symbol does not exist in the current working tree, the |
| 184 | historical version is appended to the end of the file. |
| 185 | |
| 186 | The file path component of ADDRESS is validated against the repo root — |
| 187 | path-traversal addresses (e.g. ``../../etc/passwd::foo``) are rejected. |
| 188 | """ |
| 189 | address: str = args.address |
| 190 | ref: str = args.ref |
| 191 | dry_run: bool = args.dry_run |
| 192 | as_json: bool = args.as_json |
| 193 | |
| 194 | root = require_repo() |
| 195 | |
| 196 | try: |
| 197 | repo_id = read_repo_id(root) |
| 198 | except (FileNotFoundError, json.JSONDecodeError, KeyError) as exc: |
| 199 | logger.error("Cannot read repo identity: %s", exc) |
| 200 | raise SystemExit(ExitCode.INTERNAL_ERROR) from exc |
| 201 | |
| 202 | branch = read_current_branch(root) |
| 203 | |
| 204 | if "::" not in address: |
| 205 | logger.error("ADDRESS must be a symbol address like 'src/billing.py::func'.") |
| 206 | raise SystemExit(ExitCode.USER_ERROR) |
| 207 | |
| 208 | file_rel, _sym_name = address.split("::", 1) |
| 209 | |
| 210 | # Validate the file path stays inside the repo root. |
| 211 | try: |
| 212 | contain_path(root, file_rel) |
| 213 | except ValueError as exc: |
| 214 | logger.error("%s", exc) |
| 215 | raise SystemExit(ExitCode.USER_ERROR) from exc |
| 216 | |
| 217 | commit = resolve_commit_ref(root, repo_id, branch, ref) |
| 218 | if commit is None: |
| 219 | logger.error("Commit %r not found.", ref) |
| 220 | raise SystemExit(ExitCode.USER_ERROR) |
| 221 | |
| 222 | # ------------------------------------------------------------------ |
| 223 | # Load the historical blob and locate the symbol inside it. |
| 224 | # ------------------------------------------------------------------ |
| 225 | manifest = get_commit_snapshot_manifest(root, commit.commit_id) or {} |
| 226 | obj_id = manifest.get(file_rel) |
| 227 | if obj_id is None: |
| 228 | logger.error("'%s' is not in snapshot %s.", file_rel, commit.commit_id[:8]) |
| 229 | raise SystemExit(ExitCode.USER_ERROR) |
| 230 | |
| 231 | historical_raw = read_object(root, obj_id) |
| 232 | if historical_raw is None: |
| 233 | logger.error("Blob %s missing from object store.", obj_id[:8]) |
| 234 | raise SystemExit(ExitCode.USER_ERROR) |
| 235 | |
| 236 | hist_rec = _find_symbol_in_source(historical_raw, file_rel, address) |
| 237 | if hist_rec is None: |
| 238 | logger.error("Symbol '%s' not found in commit %s.", address, commit.commit_id[:8]) |
| 239 | raise SystemExit(ExitCode.USER_ERROR) |
| 240 | |
| 241 | historical_lines = _extract_lines( |
| 242 | historical_raw, hist_rec["lineno"], hist_rec["end_lineno"] |
| 243 | ) |
| 244 | # Guard against a corrupted snapshot producing empty content. Writing |
| 245 | # an empty replacement would silently delete the symbol — hard to debug. |
| 246 | if not historical_lines: |
| 247 | logger.error( |
| 248 | "Symbol '%s' at commit %s produced no extractable lines — " |
| 249 | "snapshot may be corrupted. Aborting without writing.", |
| 250 | address, commit.commit_id[:8], |
| 251 | ) |
| 252 | raise SystemExit(ExitCode.INTERNAL_ERROR) |
| 253 | |
| 254 | # ------------------------------------------------------------------ |
| 255 | # Load the current working-tree file. |
| 256 | # Read once — reuse bytes for both line-list and symbol lookup. |
| 257 | # ------------------------------------------------------------------ |
| 258 | working_file = root / file_rel |
| 259 | if working_file.exists(): |
| 260 | current_raw = working_file.read_bytes() |
| 261 | current_lines = current_raw.decode("utf-8", errors="replace").splitlines(keepends=True) |
| 262 | # Pass file_rel (not the absolute path) so parse_symbols builds |
| 263 | # addresses matching the address format used throughout Muse. |
| 264 | cur_rec = _find_symbol_in_source(current_raw, file_rel, address) |
| 265 | else: |
| 266 | current_raw = b"" |
| 267 | current_lines = [] |
| 268 | cur_rec = None |
| 269 | |
| 270 | # ------------------------------------------------------------------ |
| 271 | # No-op detection: bail out early when the bodies already match. |
| 272 | # Only skip when NOT in dry-run mode — a dry-run explicitly requests |
| 273 | # "show me what would happen," so it must go through the full diff |
| 274 | # path even when the result would be unchanged. Skipping early in |
| 275 | # dry-run mode would omit verified_preview from the JSON, breaking |
| 276 | # agent pipelines that always pass --dry-run before writing. |
| 277 | # ------------------------------------------------------------------ |
| 278 | already_current = ( |
| 279 | cur_rec is not None |
| 280 | and cur_rec["body_hash"] == hist_rec["body_hash"] |
| 281 | ) |
| 282 | if already_current and not dry_run: |
| 283 | assert cur_rec is not None # narrowed by already_current check above |
| 284 | if as_json: |
| 285 | print(json.dumps({ |
| 286 | "schema_version": __version__, |
| 287 | "address": address, |
| 288 | "file": file_rel, |
| 289 | "branch": branch, |
| 290 | "restored_from": commit.commit_id[:8], |
| 291 | "dry_run": False, |
| 292 | "changed": False, |
| 293 | "appended": False, |
| 294 | "current_start": cur_rec["lineno"], |
| 295 | "current_end": cur_rec["end_lineno"], |
| 296 | "historical_line_count": len(historical_lines), |
| 297 | "diff_lines": [], |
| 298 | "verified": True, |
| 299 | }, indent=2)) |
| 300 | else: |
| 301 | print( |
| 302 | f"✅ {address} already matches commit {commit.commit_id[:8]}" |
| 303 | " — nothing to do." |
| 304 | ) |
| 305 | return |
| 306 | |
| 307 | # ------------------------------------------------------------------ |
| 308 | # Compute the patched file content. |
| 309 | # ------------------------------------------------------------------ |
| 310 | appended = cur_rec is None |
| 311 | |
| 312 | if cur_rec is not None: |
| 313 | cur_start, cur_end = cur_rec["lineno"], cur_rec["end_lineno"] |
| 314 | new_lines = current_lines[:cur_start - 1] + historical_lines + current_lines[cur_end:] |
| 315 | else: |
| 316 | cur_start = cur_end = 0 |
| 317 | new_lines = current_lines + ["\n"] + historical_lines |
| 318 | |
| 319 | # ------------------------------------------------------------------ |
| 320 | # Dry run — report without writing. |
| 321 | # ------------------------------------------------------------------ |
| 322 | if dry_run: |
| 323 | diff_lines = list(difflib.unified_diff( |
| 324 | current_lines, |
| 325 | new_lines, |
| 326 | fromfile="current", |
| 327 | tofile="historical", |
| 328 | lineterm="", |
| 329 | )) |
| 330 | # Preview whether the resulting file would be parseable — lets agents |
| 331 | # detect malformed output before committing to a write. |
| 332 | verified_preview = _find_symbol_in_source( |
| 333 | "".join(new_lines).encode("utf-8"), file_rel, address |
| 334 | ) is not None |
| 335 | if as_json: |
| 336 | print(json.dumps({ |
| 337 | "schema_version": __version__, |
| 338 | "address": address, |
| 339 | "file": file_rel, |
| 340 | "branch": branch, |
| 341 | "restored_from": commit.commit_id[:8], |
| 342 | "dry_run": True, |
| 343 | "changed": not already_current, |
| 344 | "appended": appended, |
| 345 | "current_start": cur_start, |
| 346 | "current_end": cur_end, |
| 347 | "historical_line_count": len(historical_lines), |
| 348 | "diff_lines": diff_lines, |
| 349 | "verified_preview": verified_preview, |
| 350 | }, indent=2)) |
| 351 | else: |
| 352 | print("Dry run — no files will be written.\n") |
| 353 | print(f"Restoring: {sanitize_display(address)}") |
| 354 | print(f" from commit: {commit.commit_id[:8]} ({commit.committed_at.date()})") |
| 355 | if not verified_preview: |
| 356 | print(" ⚠️ Warning: result would not be parseable at this address.") |
| 357 | print("\n" + "".join(diff_lines)) |
| 358 | return |
| 359 | |
| 360 | # ------------------------------------------------------------------ |
| 361 | # Write the patched file. |
| 362 | # ------------------------------------------------------------------ |
| 363 | written_content = "".join(new_lines) |
| 364 | working_file.write_text(written_content, encoding="utf-8") |
| 365 | |
| 366 | # Post-write verification: confirm the symbol is parseable at its address |
| 367 | # in the content we just wrote. Uses the in-memory string to avoid a |
| 368 | # second disk read. Catches splice edge-cases (encoding artifacts, AST |
| 369 | # parse failures) that would silently produce a broken file. |
| 370 | verified = _find_symbol_in_source( |
| 371 | written_content.encode("utf-8"), file_rel, address |
| 372 | ) is not None |
| 373 | if not verified: |
| 374 | logger.warning( |
| 375 | "Post-write verification failed: '%s' not found after restore. " |
| 376 | "The file was written but the symbol may not be at the expected " |
| 377 | "address. Inspect with: muse code symbols --file %s", |
| 378 | address, file_rel, |
| 379 | ) |
| 380 | |
| 381 | if as_json: |
| 382 | print(json.dumps({ |
| 383 | "schema_version": __version__, |
| 384 | "address": address, |
| 385 | "file": file_rel, |
| 386 | "branch": branch, |
| 387 | "restored_from": commit.commit_id[:8], |
| 388 | "dry_run": False, |
| 389 | "changed": True, |
| 390 | "appended": appended, |
| 391 | "current_start": cur_start, |
| 392 | "current_end": cur_end, |
| 393 | "historical_line_count": len(historical_lines), |
| 394 | "diff_lines": [], |
| 395 | "verified": verified, |
| 396 | }, indent=2)) |
| 397 | else: |
| 398 | print(f"Restoring: {sanitize_display(address)}") |
| 399 | print(f" from commit: {commit.commit_id[:8]} ({commit.committed_at.date()})") |
| 400 | if not appended: |
| 401 | print( |
| 402 | f" lines {cur_start}–{cur_end} → replaced with " |
| 403 | f"{len(historical_lines)} historical line(s)" |
| 404 | ) |
| 405 | else: |
| 406 | print(" symbol not found in working tree — appending at end of file") |
| 407 | if verified: |
| 408 | print(f"✅ Written to {file_rel}") |
| 409 | else: |
| 410 | print(f"⚠️ Written to {file_rel} — post-write verification failed") |