git: implement merge/checkout via pygit2 (#5245)

* git: implement merge (gitpython, pygit2)

* pylint: add pygit2 as C extension

* git: s/checkout_paths/checkout_index/ and add pygit implementation

* git: support basic checkout/merge conflict strategies in checkout_index

* fix relpath issue in pygit merge

* git: add stubs for add --update

* experiments: drop remaining gitpython specific scm usage

.pylintrc

@@ -3,7 +3,7 @@
 # A comma-separated list of package or module names from where C extensions may
 # be loaded. Extensions are loading into the active Python interpreter and may
 # run arbitrary code.
-extension-pkg-whitelist=
+extension-pkg-whitelist=pygit2
 
 # Add files or directories to the blacklist. They should be base names, not
 # paths.

dvc/repo/experiments/__init__.py

@@ -251,7 +251,7 @@ class Experiments:
         ]
         if lock_files:
             self.scm.reset(paths=lock_files)
-            self.scm.checkout_paths(lock_files, force=True)
+            self.scm.checkout_index(paths=lock_files, force=True)
 
     def _stash_msg(
         self,

dvc/repo/experiments/apply.py

@@ -20,9 +20,8 @@ logger = logging.getLogger(__name__)
 @locked
 @scm_context
 def apply(repo, rev, force=False, **kwargs):
-    from git.exc import GitCommandError
-
     from dvc.repo.checkout import checkout as dvc_checkout
+    from dvc.scm.base import SCMError
 
     exps = repo.experiments
 
@@ -44,23 +43,23 @@ def apply(repo, rev, force=False, **kwargs):
     else:
         workspace = None
 
-    repo.scm.gitpython.repo.git.merge(exp_rev, squash=True, no_commit=True)
+    repo.scm.merge(exp_rev, commit=False)
 
     if workspace:
         try:
             repo.scm.stash.apply(workspace)
-        except GitCommandError:
+        except SCMError:
             # Applied experiment conflicts with user's workspace changes
             if force:
                 # prefer applied experiment changes over prior stashed changes
-                repo.scm.gitpython.repo.git.checkout("--ours", "--", ".")
+                repo.scm.checkout_index(ours=True)
             else:
                 # revert applied changes and restore user's workspace
                 repo.scm.reset(hard=True)
                 repo.scm.stash.pop()
                 raise ApplyConflictError(rev)
         repo.scm.stash.drop()
-    repo.scm.gitpython.repo.git.reset()
+    repo.scm.reset()
 
     if stash_rev:
         args_path = os.path.join(repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE)

dvc/repo/experiments/executor/base.py

@@ -106,9 +106,8 @@ class BaseExecutor(ABC):
             # committing
             head = EXEC_BRANCH if branch else EXEC_HEAD
             self.scm.checkout(head, detach=True)
-            self.scm.gitpython.repo.git.merge(
-                EXEC_MERGE, squash=True, no_commit=True
-            )
+            merge_rev = self.scm.get_ref(EXEC_MERGE)
+            self.scm.merge(merge_rev, squash=True, commit=False)
         finally:
             os.chdir(cwd)
 
@@ -400,7 +399,7 @@ class BaseExecutor(ABC):
                 raise ExperimentExistsError(ref_info.name)
             logger.debug("Commit to new experiment branch '%s'", branch)
 
-        scm.gitpython.repo.git.add(update=True)
+        scm.add([], update=True)
         scm.commit(f"dvc: commit experiment {exp_hash}", no_verify=True)
         new_rev = scm.get_rev()
         scm.set_ref(branch, new_rev, old_ref=old_ref)

dvc/scm/base.py

@@ -34,6 +34,10 @@ class NoSCMError(SCMError):
         super().__init__(msg)
 
 
+class MergeConflictError(SCMError):
+    pass
+
+
 class Base:
     """Base class for source control management driver implementations."""
 

dvc/scm/git/__init__.py

@@ -385,8 +385,9 @@ class Git(Base):
     describe = partialmethod(_backend_func, "describe")
     diff = partialmethod(_backend_func, "diff")
     reset = partialmethod(_backend_func, "reset")
-    checkout_paths = partialmethod(_backend_func, "checkout_paths")
+    checkout_index = partialmethod(_backend_func, "checkout_index")
     status = partialmethod(_backend_func, "status")
+    merge = partialmethod(_backend_func, "merge")
 
     def resolve_rev(self, rev: str) -> str:
         from dvc.repo.experiments.utils import exp_refs_by_name

dvc/scm/git/backend/base.py

@@ -52,7 +52,7 @@ class BaseGitBackend(ABC):
         pass
 
     @abstractmethod
-    def add(self, paths: Iterable[str]):
+    def add(self, paths: Iterable[str], update=False):
         pass
 
     @abstractmethod
@@ -290,8 +290,14 @@ class BaseGitBackend(ABC):
         """Reset current git HEAD."""
 
     @abstractmethod
-    def checkout_paths(self, paths: Iterable[str], force: bool = False):
-        """Checkout the specified paths from HEAD index."""
+    def checkout_index(
+        self,
+        paths: Optional[Iterable[str]] = None,
+        force: bool = False,
+        ours: bool = False,
+        theirs: bool = False,
+    ):
+        """Checkout the specified paths from index."""
 
     @abstractmethod
     def status(
@@ -308,3 +314,30 @@ class BaseGitBackend(ABC):
 
     def _reset(self) -> None:
         pass
+
+    @abstractmethod
+    def merge(
+        self,
+        rev: str,
+        commit: bool = True,
+        msg: Optional[str] = None,
+        squash: bool = False,
+    ) -> Optional[str]:
+        """Merge specified commit into the current (working copy) branch.
+
+        Args:
+            rev: Git revision to merge.
+            commit: If True, merge will be committed using `msg` as the commit
+                message. If False, the merge changes will be staged but not
+                committed.
+            msg: Merge commit message, required if `commit` is True.
+            squash: If True, working tree state will be updated with merge
+                results, but no commit will be made and no changes will be
+                staged. `commit` must be False when `squash` is True.
+
+        If conflicts are present after merging, MergeConflictError will be
+        raised. The caller is responsible for either resolving the conflicts
+        or resetting the repository state.
+
+        Returns revision of the merge commit or None if no commit was made.
+        """

dvc/scm/git/backend/dulwich.py

@@ -117,9 +117,12 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
     def dir(self) -> str:
         return self.repo.commondir()
 
-    def add(self, paths: Iterable[str]):
+    def add(self, paths: Iterable[str], update=False):
         from dvc.utils.fs import walk_files
 
+        if update:
+            raise NotImplementedError
+
         if isinstance(paths, str):
             paths = [paths]
 
@@ -539,7 +542,13 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
     def reset(self, hard: bool = False, paths: Iterable[str] = None):
         raise NotImplementedError
 
-    def checkout_paths(self, paths: Iterable[str], force: bool = False):
+    def checkout_index(
+        self,
+        paths: Optional[Iterable[str]] = None,
+        force: bool = False,
+        ours: bool = False,
+        theirs: bool = False,
+    ):
         raise NotImplementedError
 
     def status(
@@ -562,3 +571,12 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
 
     def _reset(self) -> None:
         self.__dict__.pop("ignore_manager", None)
+
+    def merge(
+        self,
+        rev: str,
+        commit: bool = True,
+        msg: Optional[str] = None,
+        squash: bool = False,
+    ) -> Optional[str]:
+        raise NotImplementedError

dvc/scm/git/backend/gitpython.py

@@ -8,7 +8,7 @@ from typing import Callable, Iterable, Mapping, Optional, Tuple
 from funcy import first, ignore
 
 from dvc.progress import Tqdm
-from dvc.scm.base import CloneError, RevError, SCMError
+from dvc.scm.base import CloneError, MergeConflictError, RevError, SCMError
 from dvc.utils import fix_env, is_binary
 
 from ..objects import GitObject
@@ -188,11 +188,14 @@ class GitPythonBackend(BaseGitBackend):  # pylint:disable=abstract-method
     def dir(self) -> str:
         return self.repo.git_dir
 
-    def add(self, paths: Iterable[str]):
+    def add(self, paths: Iterable[str], update=False):
         # NOTE: GitPython is not currently able to handle index version >= 3.
         # See https://github.com/iterative/dvc/issues/610 for more details.
         try:
-            self.repo.index.add(paths)
+            if update:
+                self.git.add(*paths, update=True)
+            else:
+                self.repo.index.add(paths)
         except AssertionError:
             msg = (
                 "failed to add '{}' to git. You can add those files "
@@ -487,7 +490,16 @@ class GitPythonBackend(BaseGitBackend):  # pylint:disable=abstract-method
         return commit.hexsha, False
 
     def _stash_apply(self, rev: str):
-        self.git.stash("apply", rev)
+        from git.exc import GitCommandError
+
+        try:
+            self.git.stash("apply", rev)
+        except GitCommandError as exc:
+            if "CONFLICT" in str(exc):
+                raise MergeConflictError(
+                    "Stash apply resulted in merge conflicts"
+                ) from exc
+            raise SCMError("Could not apply stash") from exc
 
     def reflog_delete(
         self, ref: str, updateref: bool = False, rewrite: bool = False
@@ -515,12 +527,56 @@ class GitPythonBackend(BaseGitBackend):  # pylint:disable=abstract-method
     def reset(self, hard: bool = False, paths: Iterable[str] = None):
         self.repo.head.reset(index=True, working_tree=hard, paths=paths)
 
-    def checkout_paths(self, paths: Iterable[str], force: bool = False):
+    def checkout_index(
+        self,
+        paths: Optional[Iterable[str]] = None,
+        force: bool = False,
+        ours: bool = False,
+        theirs: bool = False,
+    ):
         """Checkout the specified paths from HEAD index."""
-        if paths:
+        assert not (ours and theirs)
+        if ours or theirs:
+            args = ["--ours"] if ours else ["--theirs"]
+            if force:
+                args.append("--force")
+            args.append("--")
+            if paths:
+                args.extend(list(paths))
+            else:
+                args.append(".")
+            self.repo.git.checkout(*args)
+        else:
             self.repo.index.checkout(paths=paths, force=force)
 
     def status(
         self, ignored: bool = False
     ) -> Tuple[Mapping[str, Iterable[str]], Iterable[str], Iterable[str]]:
         raise NotImplementedError
+
+    def merge(
+        self,
+        rev: str,
+        commit: bool = True,
+        msg: Optional[str] = None,
+        squash: bool = False,
+    ) -> Optional[str]:
+        from git.exc import GitCommandError
+
+        if commit and squash:
+            raise SCMError("Cannot merge with 'squash' and 'commit'")
+
+        if commit and not msg:
+            raise SCMError("Merge commit message is required")
+
+        merge = partial(self.git.merge, rev)
+        try:
+            if commit:
+                merge(m=msg)
+                return self.get_rev()
+            merge(no_commit=True, squash=True)
+        except GitCommandError as exc:
+            if "CONFLICT" in str(exc):
+                raise MergeConflictError("Merge contained conflicts") from exc
+            raise SCMError("Merge failed") from exc
+        return None

dvc/scm/git/backend/pygit2.py

@@ -2,9 +2,9 @@ import locale
 import logging
 import os
 from io import BytesIO, StringIO
-from typing import Callable, Iterable, Mapping, Optional, Tuple
+from typing import Callable, Iterable, List, Mapping, Optional, Tuple
 
-from dvc.scm.base import SCMError
+from dvc.scm.base import MergeConflictError, SCMError
 from dvc.utils import relpath
 
 from ..objects import GitObject
@@ -84,7 +84,7 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
     def dir(self) -> str:
         raise NotImplementedError
 
-    def add(self, paths: Iterable[str]):
+    def add(self, paths: Iterable[str], update=False):
         raise NotImplementedError
 
     def commit(self, msg: str, no_verify: bool = False):
@@ -225,8 +225,55 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
     def reset(self, hard: bool = False, paths: Iterable[str] = None):
         raise NotImplementedError
 
-    def checkout_paths(self, paths: Iterable[str], force: bool = False):
-        raise NotImplementedError
+    def checkout_index(
+        self,
+        paths: Optional[Iterable[str]] = None,
+        force: bool = False,
+        ours: bool = False,
+        theirs: bool = False,
+    ):
+        from pygit2 import (
+            GIT_CHECKOUT_ALLOW_CONFLICTS,
+            GIT_CHECKOUT_FORCE,
+            GIT_CHECKOUT_RECREATE_MISSING,
+            GIT_CHECKOUT_SAFE,
+        )
+
+        assert not (ours and theirs)
+        strategy = GIT_CHECKOUT_RECREATE_MISSING
+        if force or ours or theirs:
+            strategy |= GIT_CHECKOUT_FORCE
+        else:
+            strategy |= GIT_CHECKOUT_SAFE
+
+        if ours or theirs:
+            strategy |= GIT_CHECKOUT_ALLOW_CONFLICTS
+
+        index = self.repo.index
+        if paths:
+            path_list: Optional[List[str]] = [
+                relpath(path, self.root_dir) for path in paths
+            ]
+        else:
+            path_list = None
+        self.repo.checkout_index(
+            index=index, paths=path_list, strategy=strategy,
+        )
+
+        if index.conflicts and (ours or theirs):
+            for ancestor, ours_entry, theirs_entry in index.conflicts:
+                if not ancestor:
+                    continue
+                if ours:
+                    entry = ours_entry
+                    index.add(ours_entry)
+                else:
+                    entry = theirs_entry
+                path = os.path.join(self.root_dir, entry.path)
+                with open(path, "wb") as fobj:
+                    fobj.write(self.repo.get(entry.id).read_raw())
+                index.add(entry.path)
+            index.write()
 
     def iter_remote_refs(self, url: str, base: Optional[str] = None):
         raise NotImplementedError
@@ -235,3 +282,38 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
         self, ignored: bool = False
     ) -> Tuple[Mapping[str, Iterable[str]], Iterable[str], Iterable[str]]:
         raise NotImplementedError
+
+    def merge(
+        self,
+        rev: str,
+        commit: bool = True,
+        msg: Optional[str] = None,
+        squash: bool = False,
+    ) -> Optional[str]:
+        from pygit2 import GIT_RESET_MIXED, GitError
+
+        if commit and squash:
+            raise SCMError("Cannot merge with 'squash' and 'commit'")
+
+        if commit and not msg:
+            raise SCMError("Merge commit message is required")
+
+        try:
+            self.repo.merge(rev)
+        except GitError as exc:
+            raise SCMError("Merge failed") from exc
+
+        if self.repo.index.conflicts:
+            raise MergeConflictError("Merge contained conflicts")
+
+        if commit:
+            user = self.repo.default_signature
+            tree = self.repo.index.write_tree()
+            merge_commit = self.repo.create_commit(
+                "HEAD", user, user, msg, tree, [self.repo.head.target, rev]
+            )
+            return str(merge_commit)
+        if squash:
+            self.repo.reset(self.repo.head.target, GIT_RESET_MIXED)
+            self.repo.state_cleanup()
+        return None

tests/unit/scm/test_git.py

@@ -8,11 +8,12 @@ from tests.basic_env import TestDvcGit
 
 # Behaves the same as SCM but will test against all suported Git backends.
 # tmp_dir.scm will still contain a default SCM instance.
-@pytest.fixture(params=["gitpython", "dulwich"])
+@pytest.fixture(params=["gitpython", "dulwich", "pygit2"])
 def git(tmp_dir, scm, request):
     from dvc.scm.git import Git
 
     git_ = Git(os.fspath(tmp_dir), backends=[request.param])
+    git_.test_backend = request.param
     yield git_
     git_.close()
 
@@ -140,6 +141,9 @@ def test_branch_revs(tmp_dir, scm):
 
 
 def test_set_ref(tmp_dir, git):
+    if git.test_backend == "pygit2":
+        pytest.skip()
+
     tmp_dir.scm_gen({"file": "0"}, commit="init")
     init_rev = tmp_dir.scm.get_rev()
     tmp_dir.scm_gen({"file": "1"}, commit="commit")
@@ -167,6 +171,9 @@ def test_set_ref(tmp_dir, git):
 
 
 def test_get_ref(tmp_dir, git):
+    if git.test_backend == "pygit2":
+        pytest.skip()
+
     tmp_dir.scm_gen({"file": "0"}, commit="init")
     init_rev = tmp_dir.scm.get_rev()
     tmp_dir.gen(
@@ -185,6 +192,9 @@ def test_get_ref(tmp_dir, git):
 
 
 def test_remove_ref(tmp_dir, git):
+    if git.test_backend == "pygit2":
+        pytest.skip()
+
     tmp_dir.scm_gen({"file": "0"}, commit="init")
     init_rev = tmp_dir.scm.get_rev()
     tmp_dir.gen(os.path.join(".git", "refs", "foo", "bar"), init_rev)
@@ -282,6 +292,8 @@ def test_list_all_commits(tmp_dir, scm):
 
 
 def test_ignore_remove_empty(tmp_dir, scm, git):
+    if git.test_backend == "pygit2":
+        pytest.skip()
 
     test_entries = [
         {"entry": "/foo1", "path": f"{tmp_dir}/foo1"},
@@ -310,6 +322,9 @@ def test_ignore_remove_empty(tmp_dir, scm, git):
 def test_commit_no_verify(tmp_dir, scm, git, hook):
     import stat
 
+    if git.test_backend == "pygit2":
+        pytest.skip()
+
     hook_file = os.path.join(".git", "hooks", hook)
     tmp_dir.gen(
         hook_file, "#!/usr/bin/env python\nimport sys\nsys.exit(1)",
@@ -321,3 +336,77 @@ def test_commit_no_verify(tmp_dir, scm, git, hook):
     with pytest.raises(SCMError):
         git.commit("commit foo")
     git.commit("commit foo", no_verify=True)
+
+
+@pytest.mark.parametrize("squash", [True, False])
+def test_merge(tmp_dir, scm, git, squash):
+    from dvc.scm.base import MergeConflictError
+
+    if git.test_backend == "dulwich":
+        pytest.skip()
+
+    tmp_dir.scm_gen("foo", "foo", commit="init")
+    init_rev = scm.get_rev()
+
+    scm.checkout("branch", create_new=True)
+    tmp_dir.scm_gen("foo", "bar", commit="bar")
+    branch = scm.resolve_rev("branch")
+
+    scm.checkout("master")
+
+    with pytest.raises(MergeConflictError):
+        tmp_dir.scm_gen("foo", "baz", commit="baz")
+        git.merge(branch, commit=not squash, squash=squash, msg="merge")
+
+    scm.gitpython.git.reset(init_rev, hard=True)
+    merge_rev = git.merge(
+        branch, commit=not squash, squash=squash, msg="merge"
+    )
+    assert (tmp_dir / "foo").read_text() == "bar"
+    if squash:
+        assert merge_rev is None
+        assert scm.get_rev() == init_rev
+    else:
+        assert scm.get_rev() == merge_rev
+
+
+def test_checkout_index(tmp_dir, scm, git):
+    if git.test_backend == "dulwich":
+        pytest.skip()
+
+    tmp_dir.scm_gen({"foo": "foo", "bar": "bar"}, commit="init")
+    tmp_dir.gen("foo", "baz")
+
+    git.checkout_index(["foo"], force=True)
+    assert (tmp_dir / "foo").read_text() == "foo"
+
+    tmp_dir.gen({"foo": "baz", "bar": "baz"})
+    git.checkout_index(force=True)
+    assert (tmp_dir / "foo").read_text() == "foo"
+    assert (tmp_dir / "bar").read_text() == "bar"
+
+
+@pytest.mark.parametrize(
+    "strategy, expected", [("ours", "baz"), ("theirs", "bar")],
+)
+def test_checkout_index_conflicts(tmp_dir, scm, git, strategy, expected):
+    from dvc.scm.base import MergeConflictError
+
+    if git.test_backend == "dulwich":
+        pytest.skip()
+
+    tmp_dir.scm_gen({"file": "foo"}, commit="init")
+    scm.checkout("branch", create_new=True)
+    tmp_dir.scm_gen({"file": "bar"}, commit="bar")
+    bar = scm.get_rev()
+    scm.checkout("master")
+    tmp_dir.scm_gen({"file": "baz"}, commit="baz")
+
+    try:
+        git.merge(bar, commit=False, squash=True)
+    except MergeConflictError:
+        if strategy == "ours":
+            git.checkout_index(ours=True)
+        else:
+            git.checkout_index(theirs=True)
+    assert (tmp_dir / "file").read_text() == expected