restore.py
python
sha256:94f494a8f59e4b708ebb89e304209737ce34324a33aae83840a6f6b06d7b8d9d
docs: add domain-extensibility.md — the two-axis breadth/de…
Sonnet 5
4 days ago
| 1 | """``muse restore`` — restore working-tree files and/or stage entries. |
| 2 | |
| 3 | The focused, explicit alternative to ``muse checkout -- <file>`` for file |
| 4 | restoration. It never moves HEAD or switches branches. |
| 5 | |
| 6 | Targets |
| 7 | ------- |
| 8 | - **Working tree** (default): overwrite the on-disk file with the content |
| 9 | from the stage (if staged) or HEAD (if clean). The stage is untouched. |
| 10 | - **Stage only** (``--staged``): reset the stage entry back to HEAD state — |
| 11 | i.e. remove any staged modification, deletion, or addition — without |
| 12 | touching the on-disk file. |
| 13 | - **Both** (``--staged --worktree``): clear the stage entry *and* restore the |
| 14 | on-disk file from HEAD (or ``--source``). |
| 15 | - **Arbitrary source** (``--source <ref>``): use a commit ID or branch name |
| 16 | instead of HEAD as the restoration source. |
| 17 | |
| 18 | Behaviour per flag combination |
| 19 | -------------------------------- |
| 20 | ``muse restore file.py`` |
| 21 | Restore the working-tree file from the staged version (if staged) or |
| 22 | from HEAD. Stage is not modified. |
| 23 | |
| 24 | ``muse restore --staged file.py`` |
| 25 | Remove *file.py*'s entry from the stage so it matches HEAD: |
| 26 | - Staged as ``"M"`` → remove entry (HEAD version is back in effect) |
| 27 | - Staged as ``"D"`` → remove entry (undelete) |
| 28 | - Staged as ``"A"`` → remove entry (un-track new file; disk untouched) |
| 29 | - Not staged → no-op |
| 30 | |
| 31 | ``muse restore --staged --worktree file.py`` |
| 32 | Do both: clear the stage entry (as above) and restore the disk file from |
| 33 | HEAD (or ``--source``). |
| 34 | |
| 35 | ``muse restore --source <ref> file.py`` |
| 36 | Use *ref*'s snapshot manifest as the source instead of HEAD. Works with |
| 37 | ``--staged`` and ``--staged --worktree`` too. |
| 38 | |
| 39 | JSON schema (``--json``):: |
| 40 | |
| 41 | { |
| 42 | "restored": ["file1.py", ...], |
| 43 | "not_found": ["missing.py", ...], |
| 44 | "dry_run": true | false, |
| 45 | "staged": true | false, |
| 46 | "worktree": true | false, |
| 47 | "duration_ms": 12.3, |
| 48 | "exit_code": 0 |
| 49 | } |
| 50 | |
| 51 | Exit codes:: |
| 52 | |
| 53 | 0 — success (all paths restored, or nothing to do) |
| 54 | 1 — user error: file not in source, ref not found, path traversal |
| 55 | 2 — not a Muse repository |
| 56 | 3 — I/O error or object missing from store |
| 57 | |
| 58 | Examples:: |
| 59 | |
| 60 | muse restore a.py # discard working-tree changes |
| 61 | muse restore --staged a.py # unstage a.py |
| 62 | muse restore --staged --worktree a.py # full reset: stage + disk |
| 63 | muse restore --source feat a.py # restore from branch 'feat' |
| 64 | muse restore --source abc123 a.py b.py # restore two files from commit |
| 65 | muse restore --dry-run --json a.py # preview in JSON |
| 66 | """ |
| 67 | |
| 68 | import argparse |
| 69 | import json as _json |
| 70 | import logging |
| 71 | import pathlib |
| 72 | import sys |
| 73 | from muse.core.envelope import EnvelopeJson, make_envelope |
| 74 | from muse.core.errors import ExitCode |
| 75 | from muse.core.object_store import restore_object |
| 76 | from muse.core.repo import require_repo |
| 77 | from muse.core.types import Manifest |
| 78 | from muse.core.refs import ( |
| 79 | get_head_commit_id, |
| 80 | read_current_branch, |
| 81 | ) |
| 82 | from muse.core.commits import ( |
| 83 | read_commit, |
| 84 | resolve_commit_ref, |
| 85 | ) |
| 86 | from muse.core.snapshots import read_snapshot |
| 87 | from muse.core.validation import contain_path, sanitize_display |
| 88 | from muse.core.timing import start_timer |
| 89 | from muse.plugins.code.stage import StagedFileMap, read_stage, write_stage |
| 90 | |
| 91 | logger = logging.getLogger(__name__) |
| 92 | |
| 93 | # --------------------------------------------------------------------------- |
| 94 | # JSON output type |
| 95 | # --------------------------------------------------------------------------- |
| 96 | |
| 97 | class _RestoreResult(EnvelopeJson): |
| 98 | """Machine-readable output for ``muse restore --json``.""" |
| 99 | |
| 100 | restored: list[str] |
| 101 | not_found: list[str] |
| 102 | dry_run: bool |
| 103 | staged: bool |
| 104 | worktree: bool |
| 105 | |
| 106 | # --------------------------------------------------------------------------- |
| 107 | # Internal helpers |
| 108 | # --------------------------------------------------------------------------- |
| 109 | |
| 110 | def _resolve_source_manifest( |
| 111 | root: pathlib.Path, |
| 112 | source_ref: str | None, |
| 113 | ) -> Manifest: |
| 114 | """Return the manifest for *source_ref*, or HEAD if ``None``. |
| 115 | |
| 116 | Returns an empty dict when the repository has no commits yet, when |
| 117 | *source_ref* is ``None`` and HEAD is empty, or when *source_ref* cannot |
| 118 | be resolved to a known commit. Never raises. |
| 119 | |
| 120 | Args: |
| 121 | root: Absolute repo root. |
| 122 | source_ref: A branch name, commit ID, or ``None`` for HEAD. |
| 123 | |
| 124 | Returns: |
| 125 | Dict mapping repo-relative POSIX paths to object IDs, or ``{}`` when |
| 126 | no source is available. |
| 127 | """ |
| 128 | try: |
| 129 | branch = read_current_branch(root) |
| 130 | |
| 131 | if source_ref is None: |
| 132 | commit_id = get_head_commit_id(root, branch) |
| 133 | if not commit_id: |
| 134 | logger.debug("_resolve_source_manifest: repo has no commits") |
| 135 | return {} |
| 136 | commit = read_commit(root, commit_id) |
| 137 | else: |
| 138 | commit = resolve_commit_ref(root, branch, source_ref) |
| 139 | |
| 140 | if commit is None: |
| 141 | logger.debug("_resolve_source_manifest: ref %r resolved to None", source_ref) |
| 142 | return {} |
| 143 | snap = read_snapshot(root, commit.snapshot_id) |
| 144 | return dict(snap.manifest) if snap else {} |
| 145 | except Exception: |
| 146 | logger.debug("_resolve_source_manifest: exception resolving ref %r", source_ref, exc_info=True) |
| 147 | return {} |
| 148 | |
| 149 | def _resolve_file_path(root: pathlib.Path, raw: str) -> str: |
| 150 | """Resolve *raw* to a repo-relative POSIX path, rejecting path traversal. |
| 151 | |
| 152 | Handles both relative paths (resolved against CWD first, then repo root) |
| 153 | and absolute paths. Rejects any path that resolves outside *root*. |
| 154 | |
| 155 | Args: |
| 156 | root: Absolute repo root. |
| 157 | raw: Raw path as given by the user. |
| 158 | |
| 159 | Returns: |
| 160 | POSIX-style path relative to *root* (e.g. ``"src/auth.py"``). |
| 161 | |
| 162 | Raises: |
| 163 | SystemExit(USER_ERROR): path escapes the repository root. |
| 164 | """ |
| 165 | p = pathlib.Path(raw) |
| 166 | if not p.is_absolute(): |
| 167 | cwd_candidate = (pathlib.Path.cwd() / p).resolve() |
| 168 | try: |
| 169 | cwd_candidate.relative_to(root.resolve()) |
| 170 | abs_target = cwd_candidate |
| 171 | except ValueError: |
| 172 | abs_target = (root / p).resolve() |
| 173 | else: |
| 174 | abs_target = p.resolve() |
| 175 | |
| 176 | try: |
| 177 | rel = abs_target.relative_to(root.resolve()) |
| 178 | except ValueError: |
| 179 | print( |
| 180 | f"❌ fatal: '{sanitize_display(raw)}' is outside the repository root.", |
| 181 | file=sys.stderr, |
| 182 | ) |
| 183 | raise SystemExit(ExitCode.USER_ERROR) |
| 184 | |
| 185 | return rel.as_posix() |
| 186 | |
| 187 | # --------------------------------------------------------------------------- |
| 188 | # Registration |
| 189 | # --------------------------------------------------------------------------- |
| 190 | |
| 191 | def register( |
| 192 | subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]", |
| 193 | ) -> None: |
| 194 | """Register the ``muse restore`` subcommand.""" |
| 195 | parser = subparsers.add_parser( |
| 196 | "restore", |
| 197 | help="Restore working-tree files and/or stage entries from HEAD or a ref.", |
| 198 | description=__doc__, |
| 199 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 200 | ) |
| 201 | parser.add_argument( |
| 202 | "paths", |
| 203 | nargs="+", |
| 204 | metavar="PATH", |
| 205 | help="File(s) to restore.", |
| 206 | ) |
| 207 | parser.add_argument( |
| 208 | "--staged", "-S", |
| 209 | action="store_true", |
| 210 | dest="staged", |
| 211 | help=( |
| 212 | "Restore the stage entry from HEAD (or --source), without touching " |
| 213 | "the working-tree file. Removes staged modifications, deletions, " |
| 214 | "and new-file additions." |
| 215 | ), |
| 216 | ) |
| 217 | parser.add_argument( |
| 218 | "--worktree", "-W", |
| 219 | action="store_true", |
| 220 | dest="worktree", |
| 221 | help=( |
| 222 | "Restore the working-tree file. This is the default when " |
| 223 | "--staged is not given. Use together with --staged to reset both." |
| 224 | ), |
| 225 | ) |
| 226 | parser.add_argument( |
| 227 | "--source", "-s", |
| 228 | metavar="REF", |
| 229 | default=None, |
| 230 | dest="source", |
| 231 | help=( |
| 232 | "Use this commit ID or branch name as the restore source instead " |
| 233 | "of HEAD. Works with --staged and --worktree." |
| 234 | ), |
| 235 | ) |
| 236 | parser.add_argument( |
| 237 | "-n", "--dry-run", |
| 238 | action="store_true", |
| 239 | dest="dry_run", |
| 240 | help="Preview what would be restored without writing anything.", |
| 241 | ) |
| 242 | parser.add_argument( |
| 243 | "--json", "-j", |
| 244 | action="store_true", |
| 245 | dest="json_out", |
| 246 | help="Emit machine-readable JSON on stdout.", |
| 247 | ) |
| 248 | parser.set_defaults(func=run) |
| 249 | |
| 250 | # --------------------------------------------------------------------------- |
| 251 | # Run |
| 252 | # --------------------------------------------------------------------------- |
| 253 | |
| 254 | def run(args: argparse.Namespace) -> None: |
| 255 | """Restore files in the working tree and/or stage. |
| 256 | |
| 257 | For each path: |
| 258 | |
| 259 | 1. Resolve and validate the path (path-traversal guard). |
| 260 | 2. If ``--staged``: remove the stage entry (so it matches HEAD/source). |
| 261 | 3. If ``--worktree`` (or default): overwrite the disk file from the |
| 262 | stage entry (if present) or source manifest. |
| 263 | |
| 264 | The two actions are independent: ``--staged --worktree`` does both. |
| 265 | When neither flag is given, ``--worktree`` is the implicit default. |
| 266 | |
| 267 | Agent quickstart:: |
| 268 | |
| 269 | muse restore a.py --json |
| 270 | muse restore --staged a.py --json |
| 271 | muse restore --staged --worktree a.py --json |
| 272 | muse restore --source HEAD~3 a.py b.py --json |
| 273 | |
| 274 | JSON fields:: |
| 275 | |
| 276 | restored list[str] Repo-relative paths successfully restored |
| 277 | not_found list[str] Paths absent from the source or object store |
| 278 | dry_run bool True when no writes were made |
| 279 | staged bool True when --staged was in effect |
| 280 | worktree bool True when working-tree restore was in effect |
| 281 | |
| 282 | Exit codes:: |
| 283 | |
| 284 | 0 All paths restored (or would be in dry-run). |
| 285 | 1 User error: file not in source, ref not found, or path traversal. |
| 286 | 2 Not inside a Muse repository. |
| 287 | 3 Object missing from store (I/O or data-integrity failure). |
| 288 | """ |
| 289 | elapsed = start_timer() |
| 290 | |
| 291 | raw_paths: list[str] = args.paths |
| 292 | do_staged: bool = args.staged |
| 293 | do_worktree: bool = args.worktree |
| 294 | source_ref: str | None = args.source |
| 295 | dry_run: bool = args.dry_run |
| 296 | json_out: bool = args.json_out |
| 297 | |
| 298 | # Default: restore worktree when neither flag is given. |
| 299 | if not do_staged and not do_worktree: |
| 300 | do_worktree = True |
| 301 | |
| 302 | logger.debug( |
| 303 | "restore: paths=%r staged=%s worktree=%s source=%r dry_run=%s", |
| 304 | raw_paths, do_staged, do_worktree, source_ref, dry_run, |
| 305 | ) |
| 306 | |
| 307 | root = require_repo() |
| 308 | |
| 309 | # ── Validate source ref first ───────────────────────────────────────────── |
| 310 | if source_ref is not None: |
| 311 | # Attempt to resolve; surface a clean error for bad refs. |
| 312 | try: |
| 313 | branch = read_current_branch(root) |
| 314 | commit = resolve_commit_ref(root, branch, source_ref) |
| 315 | except Exception: |
| 316 | commit = None |
| 317 | if commit is None: |
| 318 | print( |
| 319 | f"❌ '{sanitize_display(source_ref)}' is not a known branch or commit ID.", |
| 320 | file=sys.stderr, |
| 321 | ) |
| 322 | raise SystemExit(ExitCode.USER_ERROR) |
| 323 | |
| 324 | # ── Load source manifest ────────────────────────────────────────────────── |
| 325 | source_manifest = _resolve_source_manifest(root, source_ref) |
| 326 | logger.debug("restore: source manifest has %d entries", len(source_manifest)) |
| 327 | |
| 328 | # ── Load current stage ──────────────────────────────────────────────────── |
| 329 | current_stage = read_stage(root) |
| 330 | new_stage = dict(current_stage) |
| 331 | |
| 332 | # ── Process paths ───────────────────────────────────────────────────────── |
| 333 | restored: list[str] = [] |
| 334 | not_found: list[str] = [] |
| 335 | any_user_error = False # user mistake: file absent from source, path traversal |
| 336 | any_io_error = False # infrastructure failure: blob missing from store |
| 337 | |
| 338 | for raw in raw_paths: |
| 339 | # Path validation — rejects traversal. |
| 340 | try: |
| 341 | rel = _resolve_file_path(root, raw) |
| 342 | except SystemExit: |
| 343 | # _resolve_file_path already printed to stderr; in JSON mode we |
| 344 | # rely on the not_found list + exit_code instead. |
| 345 | any_user_error = True |
| 346 | not_found.append(raw) |
| 347 | continue |
| 348 | |
| 349 | # ── --staged: reset stage entry to match source ─────────────────── |
| 350 | if do_staged: |
| 351 | staged_entry = current_stage.get(rel) |
| 352 | if staged_entry is not None: |
| 353 | # Remove the stage entry regardless of mode (M, D, or A). |
| 354 | # After removal the stage agrees with the source (HEAD by default). |
| 355 | if not dry_run: |
| 356 | new_stage.pop(rel, None) |
| 357 | logger.debug("restore: cleared stage entry for %r", rel) |
| 358 | # If not staged: no-op (already matches HEAD). |
| 359 | |
| 360 | # ── --worktree: restore the disk file ──────────────────────────── |
| 361 | if do_worktree: |
| 362 | # Determine the object_id to restore from. |
| 363 | # Priority: stage entry (if present and not being cleared by --staged) |
| 364 | # then source manifest. |
| 365 | object_id: str | None = None |
| 366 | |
| 367 | if do_staged: |
| 368 | # We just cleared the stage — restore from source manifest. |
| 369 | object_id = source_manifest.get(rel) |
| 370 | else: |
| 371 | staged_entry = current_stage.get(rel) |
| 372 | if staged_entry is not None and staged_entry["mode"] not in ("D",): |
| 373 | object_id = staged_entry["object_id"] |
| 374 | else: |
| 375 | object_id = source_manifest.get(rel) |
| 376 | |
| 377 | if object_id is None: |
| 378 | if not json_out: |
| 379 | print( |
| 380 | f"❌ '{sanitize_display(rel)}' is not in the source " |
| 381 | f"({'HEAD' if source_ref is None else sanitize_display(source_ref)}) " |
| 382 | f"manifest.", |
| 383 | file=sys.stderr, |
| 384 | ) |
| 385 | logger.debug("restore: %r not in source manifest", rel) |
| 386 | any_user_error = True |
| 387 | not_found.append(rel) |
| 388 | continue |
| 389 | |
| 390 | if not dry_run: |
| 391 | try: |
| 392 | dest = contain_path(root, rel) |
| 393 | except ValueError as exc: |
| 394 | if not json_out: |
| 395 | print(f"❌ Unsafe path '{sanitize_display(rel)}': {exc}", file=sys.stderr) |
| 396 | any_user_error = True |
| 397 | not_found.append(rel) |
| 398 | continue |
| 399 | |
| 400 | ok = restore_object(root, object_id, dest) |
| 401 | if not ok: |
| 402 | if not json_out: |
| 403 | print( |
| 404 | f"❌ Object for '{sanitize_display(rel)}' is missing from the " |
| 405 | f"object store — repository may be corrupt.", |
| 406 | file=sys.stderr, |
| 407 | ) |
| 408 | logger.error( |
| 409 | "restore: object %r for %r not found in store", object_id, rel |
| 410 | ) |
| 411 | any_io_error = True |
| 412 | not_found.append(rel) |
| 413 | continue |
| 414 | |
| 415 | logger.debug("restore: wrote %r from object %r", rel, object_id) |
| 416 | |
| 417 | restored.append(rel) |
| 418 | |
| 419 | # ── Check file-not-in-source for staged-only (no worktree check) ───────── |
| 420 | # For paths that only use --staged, we already handled them above (removing |
| 421 | # from stage or no-op). Non-existent source paths are not an error for |
| 422 | # --staged-only mode (removing a new-file staging is valid even when HEAD |
| 423 | # doesn't have it). |
| 424 | |
| 425 | # ── Commit stage changes ────────────────────────────────────────────────── |
| 426 | if do_staged and not dry_run and new_stage != current_stage: |
| 427 | write_stage(root, new_stage) |
| 428 | logger.debug("restore: wrote updated stage") |
| 429 | |
| 430 | # ── Determine final exit code ───────────────────────────────────────────── |
| 431 | if any_io_error: |
| 432 | final_exit_code: int = ExitCode.INTERNAL_ERROR |
| 433 | elif any_user_error: |
| 434 | final_exit_code = ExitCode.USER_ERROR |
| 435 | else: |
| 436 | final_exit_code = ExitCode.SUCCESS |
| 437 | |
| 438 | duration_ms = elapsed() |
| 439 | logger.debug( |
| 440 | "restore: done in %.1f ms — restored=%d not_found=%d exit_code=%d", |
| 441 | duration_ms, len(restored), len(not_found), final_exit_code, |
| 442 | ) |
| 443 | |
| 444 | # ── Output ─────────────────────────────────────────────────────────────── |
| 445 | if json_out: |
| 446 | result = _RestoreResult( |
| 447 | **make_envelope(elapsed, exit_code=final_exit_code), |
| 448 | restored=restored, |
| 449 | not_found=not_found, |
| 450 | dry_run=dry_run, |
| 451 | staged=do_staged, |
| 452 | worktree=do_worktree, |
| 453 | ) |
| 454 | print(_json.dumps(result)) |
| 455 | else: |
| 456 | for rel in restored: |
| 457 | verb = "[dry-run] Would restore" if dry_run else "Restored" |
| 458 | targets = [] |
| 459 | if do_staged: |
| 460 | targets.append("stage") |
| 461 | if do_worktree: |
| 462 | targets.append("worktree") |
| 463 | print(f"{verb}: {sanitize_display(rel)} ({', '.join(targets)})") |
| 464 | # Summary line |
| 465 | n_restored = len(restored) |
| 466 | if dry_run: |
| 467 | print(f"Would restore {n_restored} file(s).") |
| 468 | else: |
| 469 | print(f"Restored {n_restored} file(s).") |
| 470 | if not_found: |
| 471 | print(f"{len(not_found)} error(s).", file=sys.stderr) |
| 472 | for rel in not_found: |
| 473 | print(f" not found: {sanitize_display(rel)}", file=sys.stderr) |
| 474 | |
| 475 | if final_exit_code != ExitCode.SUCCESS: |
| 476 | raise SystemExit(final_exit_code) |
File History
1 commit
sha256:94f494a8f59e4b708ebb89e304209737ce34324a33aae83840a6f6b06d7b8d9d
docs: add domain-extensibility.md — the two-axis breadth/de…
Sonnet 5
4 days ago