Fix dulwich checkout

pcs/repo.py

@@ -16,7 +16,7 @@
 from pcs.identifiers import ComponentIdentifier, RepoIdentifier
 from pcs.registry import InMemoryRegistry, Registry
 from pcs.specs import NestedRepoSpec, RepoSpec, RepoSpecKeys, flatten_spec
-from pcs.utils import AOS_GLOBAL_REPOS_DIR, clear_cache_path
+from pcs.utils import AOS_GLOBAL_REPOS_DIR, clear_cache_path, dulwich_checkout
 
 logger = logging.getLogger(__name__)
 
@@ -361,7 +361,9 @@ def _checkout_version(self, local_repo_path: Path, version: str) -> None:
         if treeish is None:
             treeish = parse_commit(repo, to_checkout).sha().hexdigest()
 
-        porcelain.reset(repo=repo, mode="hard", treeish=treeish)
+        # Checks for a clean working directory were failing on Windows, so
+        # force the checkout since this should be a clean clone anyway.
+        dulwich_checkout(repo=repo, target=treeish, force=True)
         os.chdir(curr_dir)
 
 

pcs/utils.py

@@ -1,9 +1,12 @@
+import os
 import pprint
 import shutil
 from pathlib import Path
 from typing import Dict, Optional
 
 import yaml
+from dulwich import porcelain
+from dulwich.refs import LOCAL_BRANCH_PREFIX
 
 AOS_GLOBAL_CONFIG_DIR = Path.home() / ".agentos"
 AOS_GLOBAL_CACHE_DIR = AOS_GLOBAL_CONFIG_DIR / "cache"
@@ -163,5 +166,112 @@ def _handle_acme_r2d2(version_string):
     return _handle_agent(r2d2_path_prefix, r2d2_rename_map)
 
 
+# https://github.com/jelmer/dulwich/pull/898
+def dulwich_checkout(repo, target: bytes, force: bool = False):
+    """
+    This code is taken from the currently unmerged (but partially
+    cherry-picked) code in https://github.com/jelmer/dulwich/pull/898
+    that adds checkout() functionality to the dulwich code base.
+
+    Switch branches or restore working tree files
+    Args:
+      repo: dulwich Repo object
+      target: branch name or commit sha to checkout
+    """
+    # check repo status
+    if not force:
+        index = repo.open_index()
+        for file in porcelain.get_tree_changes(repo)["modify"]:
+            if file in index:
+                raise Exception(
+                    "trying to checkout when working directory not clean"
+                )
+
+        normalizer = repo.get_blob_normalizer()
+        filter_callback = normalizer.checkin_normalize
+
+        unstaged_changes = list(
+            porcelain.get_unstaged_changes(index, repo.path, filter_callback)
+        )
+        for file in unstaged_changes:
+            if file in index:
+                raise Exception(
+                    "Trying to checkout when working directory not clean"
+                )
+
+    current_tree = porcelain.parse_tree(repo, repo.head())
+    target_tree = porcelain.parse_tree(repo, target)
+
+    # update head
+    if (
+        target == b"HEAD"
+    ):  # do not update head while trying to checkout to HEAD
+        target = repo.head()
+    elif target in repo.refs.keys(base=LOCAL_BRANCH_PREFIX):
+        porcelain.update_head(repo, target)
+        target = repo.refs[LOCAL_BRANCH_PREFIX + target]
+    else:
+        porcelain.update_head(repo, target, detached=True)
+
+    # unstage files in the current_tree or target_tree
+    tracked_changes = []
+    for change in repo.open_index().changes_from_tree(
+        repo.object_store, target_tree.id
+    ):
+        file = (
+            change[0][0] or change[0][1]
+        )  # no matter the file is added, modified or deleted.
+        try:
+            current_entry = current_tree.lookup_path(
+                repo.object_store.__getitem__, file
+            )
+        except KeyError:
+            current_entry = None
+        try:
+            target_entry = target_tree.lookup_path(
+                repo.object_store.__getitem__, file
+            )
+        except KeyError:
+            target_entry = None
+
+        if current_entry or target_entry:
+            tracked_changes.append(file)
+    tracked_changes = [tc.decode("utf-8") for tc in tracked_changes]
+    repo.unstage(tracked_changes)
+
+    # reset tracked and unstaged file to target
+    normalizer = repo.get_blob_normalizer()
+    filter_callback = normalizer.checkin_normalize
+    unstaged_files = porcelain.get_unstaged_changes(
+        repo.open_index(), repo.path, filter_callback
+    )
+    saved_repo_path = repo.path
+    repo.path = str(repo.path)
+    for file in unstaged_files:
+        porcelain.reset_file(repo, file.decode(), b"HEAD")
+    repo.path = saved_repo_path
+
+    # remove the untracked file which in the current_file_set
+    for file in porcelain.get_untracked_paths(
+        repo.path, repo.path, repo.open_index(), exclude_ignored=True
+    ):
+        # TODO: Code below is from the original dulwich PR; had trouble
+        # getting this to work on Windows; Untracked files sitting in repo
+        # weren't being properly removed.  Went with a more direct approach.
+        #
+        # try:
+        #     current_tree.lookup_path(
+        #         repo.object_store.__getitem__, file.encode()
+        #     )
+        # except KeyError:
+        #     pass
+        # else:
+        #     os.remove(os.path.join(repo.path, file))
+
+        full_path = Path(repo.path) / Path(file)
+        if full_path.exists():
+            os.remove(full_path)
+
+
 if __name__ == "__main__":
     generate_dummy_dev_registry()

tests/test_repo.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 from pcs.component import Component
 from pcs.repo import LocalRepo, Repo
 from tests.utils import (
@@ -43,3 +45,20 @@ def test_local_to_from_registry():
     repo_from_reg = Repo.from_registry(reg, "test_id")
     assert repo.identifier == repo_from_reg.identifier
     assert repo.local_dir == repo_from_reg.local_dir
+
+
+def test_repo_checkout_bug():
+    aos_repo = Repo.from_github("agentos-project", "agentos")
+    pcs_component_path = Path("pcs") / Path("component.py")
+    # pcs/component.py exists in 4e12203
+    exists_version = "4e12203faaf84361af9432271e013ddfb927f75d"
+    exists_path = aos_repo.get_local_file_path(
+        pcs_component_path, version=exists_version
+    )
+    assert exists_path.exists()
+    # pcs/component.py does NOT exist in 07bc713
+    not_exists_version = "07bc71358b4360092b58d78f9eee6dc939e90b10"
+    not_exists_path = aos_repo.get_local_file_path(
+        pcs_component_path, version=not_exists_version
+    )
+    assert not not_exists_path.exists()