Week 4 - M2: Restore files and index on checkout

minigit/index.py

@@ -192,7 +192,78 @@ def _working_status(self) -> DiffResult:
         return result
 
     def checkout(self, tree_hash) -> None:
-        pass
+        target_entries = self.read_tree_entries(tree_hash)
+        target_by_path = {e.path: e for e in target_entries}
+
+        # Reject unsafe paths
+        blob_cache = {}
+        for entry in target_entries:
+            self._validate_checkout_path(entry.path)
+            _, data = self.store.read_object(entry.hash)
+            blob_cache[entry.path] = data
+
+        current_entries = self.read_index()
+        current_by_path = {e.path: e for e in current_entries}
+
+        # Refuse if any tracked file has local edits or is missing.
+        for entry in current_entries:
+            full_path = os.path.join(self.root, entry.path)
+            if not os.path.isfile(full_path):
+                raise MiniGitError(f"local changes would be lost: {entry.path} is missing")
+            with open(full_path, "rb") as f:
+                data = f.read()
+            disk_hash = self.store.write_object(data, "blob")
+            disk_mode = "100755" if os.access(full_path, os.X_OK) else "100644"
+            if disk_hash != entry.hash or disk_mode != entry.mode:
+                raise MiniGitError(f"local changes would be lost: {entry.path}")
+
+        # Refuse if an untracked file/directory sits where the target
+        # needs to write.
+        for entry in target_entries:
+            full_path = os.path.join(self.root, entry.path)
+            if entry.path not in current_by_path and os.path.exists(full_path):
+                raise MiniGitError(f"untracked path would be overwritten: {entry.path}")
+
+        # Remove tracked files the target doesn't have, then prune
+        for entry in current_entries:
+            if entry.path not in target_by_path:
+                full_path = os.path.join(self.root, entry.path)
+                if os.path.isfile(full_path):
+                    os.remove(full_path)
+                self._prune_empty_dirs(os.path.dirname(full_path))
+
+        # Create directories and write every target file's bytes/mode.
+        for entry in target_entries:
+            full_path = os.path.join(self.root, entry.path)
+            os.makedirs(os.path.dirname(full_path), exist_ok=True)
+            with open(full_path, "wb") as f:
+                f.write(blob_cache[entry.path])
+            os.chmod(full_path, 0o755 if entry.mode == "100755" else 0o644)
+
+        # Now matched
+        self.write_index(target_entries)
+
+    def _validate_checkout_path(self, path: str) -> None:
+        if path.startswith("/") or ".." in path.split("/"):
+            raise MiniGitError(f"unsafe path: {path}")
+        if any(part == ".minigit" for part in path.split("/")):
+            raise MiniGitError(f"path under .minigit: {path}")
+
+        parts = path.split("/")
+        current = self.root
+        for part in parts[:-1]:
+            current = os.path.join(current, part)
+            if os.path.islink(current):
+                raise MiniGitError(f"path passes through symlink: {path}")
+
+    def _prune_empty_dirs(self, dir_path: str) -> None:
+        root = os.path.realpath(self.root)
+        current = os.path.realpath(dir_path)
+        while current != root and current.startswith(root):
+            if not os.path.isdir(current) or os.listdir(current):
+                break
+            os.rmdir(current)
+            current = os.path.dirname(current)
 
 
 def cmd_add(args) -> int:

tests/test_index.py

@@ -1,6 +1,11 @@
 """Tests for minigit/index.py - WorkingTree, trees, diffing, and status."""
 
+import os
+
+import pytest
+
 from minigit.commits import CommitManager
+from minigit.errors import MiniGitError
 from minigit.index import WorkingTree, cmd_status
 from minigit.objects import ObjectStore
 
@@ -259,3 +264,147 @@ def test_status_reports_file_removed_from_index_as_untracked(tmp_path, monkeypat
     monkeypatch.chdir(tmp_path)
     cmd_status(_Args())
     assert capsys.readouterr().out == "staged:\n  file\nuntracked:\n  file\n"
+
+
+def test_checkout_restores_changed_file(tmp_path):
+    store = ObjectStore(str(tmp_path))
+    wt = WorkingTree(repo_path=str(tmp_path), store=store)
+    (tmp_path / "hello.txt").write_text("v1")
+    wt.stage_file("hello.txt")
+    tree_hash = wt.build_tree_from_index()
+
+    (tmp_path / "hello.txt").write_text("v1")  # revert to clean state first
+    wt.stage_file("hello.txt")
+
+    wt.checkout(tree_hash)
+    assert (tmp_path / "hello.txt").read_text() == "v1"
+
+
+def test_checkout_nested_paths_and_spaces(tmp_path):
+    store = ObjectStore(str(tmp_path))
+    wt = WorkingTree(repo_path=str(tmp_path), store=store)
+    (tmp_path / "src").mkdir()
+    (tmp_path / "src" / "my notes.txt").write_text("hi")
+    wt.stage_file("src/my notes.txt")
+    tree_hash = wt.build_tree_from_index()
+
+    (tmp_path / "src" / "my notes.txt").unlink()
+    (tmp_path / "src" / "my notes.txt").write_text("hi")  # re-stage clean copy
+    wt.stage_file("src/my notes.txt")
+
+    wt.checkout(tree_hash)
+    assert (tmp_path / "src" / "my notes.txt").read_text() == "hi"
+
+
+def test_checkout_restores_binary_bytes(tmp_path):
+    store = ObjectStore(str(tmp_path))
+    wt = WorkingTree(repo_path=str(tmp_path), store=store)
+    (tmp_path / "data.bin").write_bytes(b"\x00\x01\xff\xfe")
+    wt.stage_file("data.bin")
+    tree_hash = wt.build_tree_from_index()
+
+    wt.checkout(tree_hash)
+    assert (tmp_path / "data.bin").read_bytes() == b"\x00\x01\xff\xfe"
+
+
+def test_checkout_restores_executable_mode(tmp_path):
+    store = ObjectStore(str(tmp_path))
+    wt = WorkingTree(repo_path=str(tmp_path), store=store)
+    file = tmp_path / "run.sh"
+    file.write_text("echo hi")
+    file.chmod(0o755)
+    wt.stage_file("run.sh")
+    tree_hash = wt.build_tree_from_index()
+
+    file.chmod(0o644)  # revert manually
+    wt.stage_file("run.sh")  # re-stage as-is (still 644 on disk now)
+
+    wt.checkout(tree_hash)
+    assert os.access(tmp_path / "run.sh", os.X_OK)
+
+
+def test_checkout_removes_tracked_file_not_in_target(tmp_path):
+    store = ObjectStore(str(tmp_path))
+    wt = WorkingTree(repo_path=str(tmp_path), store=store)
+    (tmp_path / "keep.txt").write_text("keep")
+    (tmp_path / "gone.txt").write_text("gone")
+    wt.stage_file("keep.txt")
+    empty_target = wt.build_tree_from_index()  # tree with just keep.txt
+
+    wt.stage_file("gone.txt")  # now index has both, target only has keep.txt
+
+    wt.checkout(empty_target)
+    assert (tmp_path / "keep.txt").exists()
+    assert not (tmp_path / "gone.txt").exists()
+
+
+def test_checkout_preserves_untracked_file(tmp_path):
+    store = ObjectStore(str(tmp_path))
+    wt = WorkingTree(repo_path=str(tmp_path), store=store)
+    (tmp_path / "tracked.txt").write_text("hi")
+    wt.stage_file("tracked.txt")
+    tree_hash = wt.build_tree_from_index()
+
+    (tmp_path / "untracked.txt").write_text("leave me alone")
+
+    wt.checkout(tree_hash)
+    assert (tmp_path / "untracked.txt").read_text() == "leave me alone"
+
+
+def test_checkout_empty_tree_clears_index(tmp_path):
+    store = ObjectStore(str(tmp_path))
+    wt = WorkingTree(repo_path=str(tmp_path), store=store)
+    empty_hash = wt.build_tree_from_index()  # nothing staged -> empty tree
+
+    wt.checkout(empty_hash)
+    assert wt.read_index() == []
+
+
+def test_checkout_fails_on_local_edit_before_changing_anything(tmp_path):
+    store = ObjectStore(str(tmp_path))
+    wt = WorkingTree(repo_path=str(tmp_path), store=store)
+    file = tmp_path / "hello.txt"
+    file.write_text("original")
+    wt.stage_file("hello.txt")
+    tree_hash = wt.build_tree_from_index()
+
+    file.write_text("uncommitted edit")  # local edit, not re-staged
+
+    with pytest.raises(MiniGitError):
+        wt.checkout(tree_hash)
+    assert file.read_text() == "uncommitted edit"  # untouched
+
+
+def test_checkout_fails_on_untracked_path_blocking_target(tmp_path):
+    store = ObjectStore(str(tmp_path))
+    wt = WorkingTree(repo_path=str(tmp_path), store=store)
+    (tmp_path / "hello.txt").write_text("hi")
+    wt.stage_file("hello.txt")
+
+    wt2 = WorkingTree(repo_path=str(tmp_path), store=store)
+    (tmp_path / "new.txt").write_text("staged elsewhere")
+    wt2.stage_file("new.txt")
+    tree_with_new = wt2.build_tree_from_index()
+
+    # simulate an untracked file blocking the checkout target
+    wt3 = WorkingTree(repo_path=str(tmp_path), store=store)
+    wt3.write_index([e for e in wt3.read_index() if e.path != "new.txt"])
+
+    with pytest.raises(MiniGitError):
+        wt3.checkout(tree_with_new)
+
+
+def test_checkout_is_idempotent(tmp_path):
+    store = ObjectStore(str(tmp_path))
+    wt = WorkingTree(repo_path=str(tmp_path), store=store)
+    (tmp_path / "hello.txt").write_text("hi")
+    wt.stage_file("hello.txt")
+    tree_hash = wt.build_tree_from_index()
+
+    wt.checkout(tree_hash)
+    first_index = wt.read_index()
+    wt.checkout(tree_hash)
+    second_index = wt.read_index()
+
+    assert first_index == second_index
+    assert (tmp_path / "hello.txt").read_text() == "hi"