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