Week 4 - M2: Restore files and index on checkout

minigit/index.py

@@ -13,7 +13,7 @@
 from dataclasses import dataclass
 from typing import NamedTuple
 
-from minigit.errors import MiniGitError
+from minigit.errors import MiniGitError, ObjectCorruptError
 
 
 class IndexEntry(NamedTuple):
@@ -192,7 +192,110 @@ 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)
+            obj_type, data = self.store.read_object(entry.hash)
+            if obj_type != "blob":
+                raise ObjectCorruptError(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:
+            self._validate_checkout_path(entry.path)
+            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.hash_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}")
+
+        # Preflight every parent and destination before deleting any tracked file.
+        # Directories may become files only when all their contents are tracked.
+        removable = set(current_by_path) - set(target_by_path)
+        tracked_dirs = set()
+        for path in removable:
+            parent = os.path.dirname(path)
+            while parent:
+                tracked_dirs.add(parent)
+                parent = os.path.dirname(parent)
+
+        for entry in target_entries:
+            parts = entry.path.split("/")
+            for i in range(1, len(parts)):
+                parent = "/".join(parts[:i])
+                full_parent = os.path.join(self.root, parent)
+                if os.path.lexists(full_parent) and not os.path.isdir(full_parent):
+                    if parent not in removable:
+                        raise MiniGitError(f"untracked path would be overwritten: {parent}")
+
+            full_path = os.path.join(self.root, entry.path)
+            if os.path.isdir(full_path):
+                if entry.path not in tracked_dirs:
+                    raise MiniGitError(f"untracked path would be overwritten: {entry.path}")
+                for directory, dirnames, filenames in os.walk(full_path):
+                    for name in dirnames + filenames:
+                        child = os.path.join(directory, name)
+                        relative = os.path.relpath(child, self.root).replace(os.sep, "/")
+                        if os.path.islink(child) or relative not in removable | tracked_dirs:
+                            raise MiniGitError(f"untracked path would be overwritten: {relative}")
+            elif os.path.lexists(full_path) and entry.path not in current_by_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 (
+            os.path.isabs(path)
+            or any(part in {"", ".", ".."} for part in path.split("/"))
+            or any(char in path for char in "\\\0\r\n")
+        ):
+            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:
+            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 os.path.commonpath([root, current]) == 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_checkout_safety.py

@@ -0,0 +1,96 @@
+import pytest
+
+from minigit.errors import MiniGitError, ObjectCorruptError
+from minigit.index import IndexEntry, WorkingTree
+from minigit.objects import TreeEntry
+
+
+def target_tree(wt, path, data=b"target"):
+    obj = wt.store.write_object(data, "blob")
+    parts = path.split("/")
+    tree = wt.store.write_tree([TreeEntry("100644", "blob", obj, parts[-1])])
+    for part in reversed(parts[:-1]):
+        tree = wt.store.write_tree([TreeEntry("40000", "tree", tree, part)])
+    return tree
+
+
+def test_checkout_file_directory_transitions(tmp_path):
+    wt = WorkingTree(tmp_path)
+    flat = target_tree(wt, "path", b"file")
+    nested = target_tree(wt, "path/nested", b"nested")
+    wt.checkout(flat)
+    wt.checkout(nested)
+    assert (tmp_path / "path/nested").read_bytes() == b"nested"
+    wt.checkout(flat)
+    assert (tmp_path / "path").read_bytes() == b"file"
+    assert [e.path for e in wt.read_index()] == ["path"]
+
+
+def test_checkout_blocking_parent_fails_before_deleting_files(tmp_path):
+    wt = WorkingTree(tmp_path)
+    wt.checkout(target_tree(wt, "tracked"))
+    before = wt.read_index()
+    (tmp_path / "blocker").write_bytes(b"untracked")
+    with pytest.raises(MiniGitError):
+        wt.checkout(target_tree(wt, "blocker/nested"))
+    assert (tmp_path / "tracked").read_bytes() == b"target"
+    assert (tmp_path / "blocker").read_bytes() == b"untracked"
+    assert wt.read_index() == before
+
+
+@pytest.mark.parametrize("dangling", [False, True])
+def test_checkout_rejects_leaf_symlink(tmp_path, dangling):
+    root = tmp_path / "repo"
+    root.mkdir()
+    outside = tmp_path / "outside"
+    if not dangling:
+        outside.write_bytes(b"original")
+    wt = WorkingTree(root)
+    old = wt.store.write_object(b"original", "blob")
+    if not dangling:
+        wt.write_index([IndexEntry("100644", old, "link")])
+    (root / "link").symlink_to(outside)
+    with pytest.raises(MiniGitError):
+        wt.checkout(target_tree(wt, "link"))
+    assert (root / "link").is_symlink()
+    if dangling:
+        assert not outside.exists()
+    else:
+        assert outside.read_bytes() == b"original"
+
+
+def test_checkout_rejects_unsafe_current_index(tmp_path):
+    root = tmp_path / "repo"
+    root.mkdir()
+    outside = tmp_path / "outside"
+    outside.write_bytes(b"original")
+    wt = WorkingTree(root)
+    old = wt.store.write_object(b"original", "blob")
+    wt.write_index([IndexEntry("100644", old, "../outside")])
+    with pytest.raises(MiniGitError):
+        wt.checkout(wt.store.write_tree([]))
+    assert outside.read_bytes() == b"original"
+
+
+def test_checkout_rejects_non_blob_before_changing_files(tmp_path):
+    wt = WorkingTree(tmp_path)
+    wt.checkout(target_tree(wt, "tracked"))
+    before = wt.read_index()
+    wrong_type = wt.store.write_tree([])
+    bad_tree = wt.store.write_tree([TreeEntry("100644", "blob", wrong_type, "new")])
+    with pytest.raises(ObjectCorruptError):
+        wt.checkout(bad_tree)
+    assert (tmp_path / "tracked").read_bytes() == b"target"
+    assert wt.read_index() == before
+
+
+def test_checkout_directory_transition_preserves_untracked_children(tmp_path):
+    wt = WorkingTree(tmp_path)
+    wt.checkout(target_tree(wt, "path/tracked"))
+    (tmp_path / "path/untracked").write_bytes(b"keep")
+    before = wt.read_index()
+    with pytest.raises(MiniGitError):
+        wt.checkout(target_tree(wt, "path"))
+    assert (tmp_path / "path/tracked").read_bytes() == b"target"
+    assert (tmp_path / "path/untracked").read_bytes() == b"keep"
+    assert wt.read_index() == before

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"