pygit/dulwich: implement stash functionality (#5286)

* pygit2: disable set_ref with reflog message

* git: implement is_dirty in dulwich

* pygit2: implement _stash_apply

* pygit: update to 1.5.0

* enable set_ref(message=...) for pygit2

* git: replace backend reflog_delete in favor of _stash_drop

* pygit: support stash_push with custom refs

* update tests for backend stash_drop behavior

* pygit: raise MergeConflictError on failed stash apply

* pygit: reload index from disk before stash apply

* update stash_apply checkout strategy

* mark flaky test

dvc/scm/git/__init__.py

@@ -386,7 +386,7 @@ class Git(Base):
     _stash_iter = partialmethod(_backend_func, "_stash_iter")
     _stash_push = partialmethod(_backend_func, "_stash_push")
     _stash_apply = partialmethod(_backend_func, "_stash_apply")
-    reflog_delete = partialmethod(_backend_func, "reflog_delete")
+    _stash_drop = partialmethod(_backend_func, "_stash_drop")
     describe = partialmethod(_backend_func, "describe")
     diff = partialmethod(_backend_func, "diff")
     reset = partialmethod(_backend_func, "reset")

dvc/scm/git/backend/base.py

@@ -96,7 +96,7 @@ class BaseGitBackend(ABC):
         pass
 
     @abstractmethod
-    def is_dirty(self, **kwargs) -> bool:
+    def is_dirty(self, untracked_files: bool = False) -> bool:
         pass
 
     @abstractmethod
@@ -245,17 +245,21 @@ class BaseGitBackend(ABC):
         message: Optional[str] = None,
         include_untracked: Optional[bool] = False,
     ) -> Tuple[Optional[str], bool]:
-        """Push a commit onto the specified stash."""
+        """Push a commit onto the specified stash.
+
+        Returns a tuple of the form (rev, need_reset) where need_reset
+        indicates whether or not the workspace should be `reset --hard`
+        (some backends will not clean the workspace after creating a stash
+        commit).
+        """
 
     @abstractmethod
     def _stash_apply(self, rev: str):
         """Apply the specified stash revision."""
 
     @abstractmethod
-    def reflog_delete(
-        self, ref: str, updateref: bool = False, rewrite: bool = False
-    ):
-        """Delete the specified reflog entry."""
+    def _stash_drop(self, ref: str, index: int):
+        """Drop the specified stash revision."""
 
     @abstractmethod
     def describe(

dvc/scm/git/backend/dulwich.py

@@ -222,8 +222,9 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
                 return True
         return False
 
-    def is_dirty(self, **kwargs) -> bool:
-        raise NotImplementedError
+    def is_dirty(self, untracked_files: bool = False) -> bool:
+        staged, unstaged, untracked = self.status()
+        return bool(staged or unstaged or (untracked_files and untracked))
 
     def active_branch(self) -> str:
         raise NotImplementedError
@@ -244,7 +245,10 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
         return DulwichObject(self.repo, ".", stat.S_IFDIR, tree.id)
 
     def get_rev(self) -> str:
-        raise NotImplementedError
+        rev = self.get_ref("HEAD")
+        if rev:
+            return rev
+        raise SCMError("Empty git repo")
 
     def resolve_rev(self, rev: str) -> str:
         raise NotImplementedError
@@ -535,10 +539,17 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
     def _stash_apply(self, rev: str):
         raise NotImplementedError
 
-    def reflog_delete(
-        self, ref: str, updateref: bool = False, rewrite: bool = False
-    ):
-        raise NotImplementedError
+    def _stash_drop(self, ref: str, index: int):
+        from dvc.scm.git import Stash
+
+        if ref == Stash.DEFAULT_STASH:
+            raise NotImplementedError
+
+        stash = self._get_stash(ref)
+        try:
+            stash.drop(index)
+        except ValueError as exc:
+            raise SCMError("Failed to drop stash entry") from exc
 
     def describe(
         self,

dvc/scm/git/backend/gitpython.py

@@ -249,8 +249,8 @@ class GitPythonBackend(BaseGitBackend):  # pylint:disable=abstract-method
     def is_tracked(self, path):
         return bool(self.repo.git.ls_files(path))
 
-    def is_dirty(self, **kwargs):
-        return self.repo.is_dirty(**kwargs)
+    def is_dirty(self, untracked_files: bool = False) -> bool:
+        return self.repo.is_dirty(untracked_files=untracked_files)
 
     def active_branch(self):
         return self.repo.active_branch.name
@@ -495,16 +495,22 @@ class GitPythonBackend(BaseGitBackend):  # pylint:disable=abstract-method
                 ) from exc
             raise SCMError("Could not apply stash") from exc
 
-    def reflog_delete(
-        self, ref: str, updateref: bool = False, rewrite: bool = False
-    ):
-        args = ["delete"]
-        if updateref:
-            args.append("--updateref")
-        if rewrite:
-            args.append("--rewrite")
-        args.append(ref)
-        self.git.reflog(*args)
+    def _stash_drop(self, ref: str, index: int):
+        from git.exc import GitCommandError
+
+        from dvc.scm.git import Stash
+
+        if ref == Stash.DEFAULT_STASH:
+            self.git.stash("drop", index)
+            return
+
+        self.git.reflog(
+            "delete", "--updateref", "--rewrite", f"{ref}@{{{index}}}"
+        )
+        try:
+            self.git.reflog("exists", ref)
+        except GitCommandError:
+            self.remove_ref(ref)
 
     def describe(
         self,

dvc/scm/git/backend/pygit2.py

@@ -5,6 +5,8 @@ import stat
 from io import BytesIO, StringIO
 from typing import Callable, Iterable, List, Mapping, Optional, Tuple, Union
 
+from funcy import cached_property
+
 from dvc.scm.base import MergeConflictError, RevError, SCMError
 from dvc.utils import relpath
 
@@ -64,12 +66,20 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
         self._stashes: dict = {}
 
     def close(self):
+        if hasattr(self, "_refdb"):
+            del self._refdb
         self.repo.free()
 
     @property
     def root_dir(self) -> str:
         return self.repo.workdir
 
+    @cached_property
+    def _refdb(self):
+        from pygit2 import RefdbFsBackend
+
+        return RefdbFsBackend(self.repo)
+
     @staticmethod
     def clone(
         url: str,
@@ -136,7 +146,7 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
     def is_tracked(self, path: str) -> bool:
         raise NotImplementedError
 
-    def is_dirty(self, **kwargs) -> bool:
+    def is_dirty(self, untracked_files: bool = False) -> bool:
         raise NotImplementedError
 
     def active_branch(self) -> str:
@@ -213,12 +223,16 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
         if old_ref and old_ref != self.get_ref(name, follow=False):
             raise SCMError(f"Failed to set '{name}'")
 
-        if symbolic:
-            ref = self.repo.create_reference_symbolic(name, new_ref, True)
-        else:
-            ref = self.repo.create_reference_direct(name, new_ref, True)
         if message:
-            ref.set_target(new_ref, message)
+            self._refdb.ensure_log(name)
+        if symbolic:
+            self.repo.create_reference_symbolic(
+                name, new_ref, True, message=message
+            )
+        else:
+            self.repo.create_reference_direct(
+                name, new_ref, True, message=message
+            )
 
     def get_ref(self, name, follow: bool = True) -> Optional[str]:
         from pygit2 import GIT_REF_SYMBOLIC
@@ -305,15 +319,64 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
         message: Optional[str] = None,
         include_untracked: Optional[bool] = False,
     ) -> Tuple[Optional[str], bool]:
-        raise NotImplementedError
+        from dvc.scm.git import Stash
+
+        oid = self.repo.stash(
+            self.repo.default_signature,
+            message=message,
+            include_untracked=include_untracked,
+        )
+        commit = self.repo[oid]
+
+        if ref != Stash.DEFAULT_STASH:
+            self.set_ref(ref, commit.id, message=commit.message)
+            self.repo.stash_drop()
+        return str(oid), False
 
     def _stash_apply(self, rev: str):
-        raise NotImplementedError
+        from pygit2 import (
+            GIT_CHECKOUT_RECREATE_MISSING,
+            GIT_CHECKOUT_SAFE,
+            GitError,
+        )
 
-    def reflog_delete(
-        self, ref: str, updateref: bool = False, rewrite: bool = False
-    ):
-        raise NotImplementedError
+        from dvc.scm.git import Stash
+
+        def _apply(index):
+            try:
+                self.repo.index.read(False)
+                self.repo.stash_apply(
+                    index,
+                    strategy=GIT_CHECKOUT_SAFE | GIT_CHECKOUT_RECREATE_MISSING,
+                )
+            except GitError as exc:
+                raise MergeConflictError(
+                    "Stash apply resulted in merge conflicts"
+                ) from exc
+
+        # libgit2 stash apply only accepts refs/stash items by index. If rev is
+        # not in refs/stash, we will push it onto the stash, and then pop it
+        commit, _ref = self.repo.resolve_refish(rev)
+        stash = self.repo.references.get(Stash.DEFAULT_STASH)
+        if stash:
+            for i, entry in enumerate(stash.log()):
+                if entry.oid_new == commit.id:
+                    _apply(i)
+                    return
+
+        self.set_ref(Stash.DEFAULT_STASH, commit.id, message=commit.message)
+        try:
+            _apply(0)
+        finally:
+            self.repo.stash_drop()
+
+    def _stash_drop(self, ref: str, index: int):
+        from dvc.scm.git import Stash
+
+        if ref != Stash.DEFAULT_STASH:
+            raise NotImplementedError
+
+        self.repo.stash_drop(index)
 
     def describe(
         self,

dvc/scm/git/stash.py

@@ -1,11 +1,9 @@
 """git stash convenience wrapper."""
 
 import logging
-import os
 from typing import Optional
 
 from dvc.scm.base import SCMError
-from dvc.utils.fs import remove
 
 logger = logging.getLogger(__name__)
 
@@ -62,19 +60,12 @@ class Stash:
         self.scm._stash_apply(rev)  # pylint: disable=protected-access
 
     def drop(self, index: int = 0):
-        ref = "{0}@{{{1}}}".format(self.ref, index)
         if index < 0 or index >= len(self):
-            raise SCMError(f"Invalid stash ref '{ref}'")
-        logger.debug("Dropping '%s'", ref)
-        self.scm.reflog_delete(ref, updateref=True, rewrite=True)
-
-        # if we removed the last reflog entry, delete the ref and reflog
-        if len(self) == 0:
-            self.scm.remove_ref(self.ref)
-            parts = self.ref.split("/")
-            reflog = os.path.join(self.scm.root_dir, ".git", "logs", *parts)
-            if os.path.exists(reflog):
-                remove(reflog)
+            raise SCMError(f"Invalid stash ref '{self.ref}@{{{index}}}'")
+        logger.debug("Dropping '%s@{%d}'", self.ref, index)
+        self.scm._stash_drop(  # pylint: disable=protected-access
+            self.ref, index
+        )
 
     def clear(self):
         logger.debug("Clear stash '%s'", self.ref)

setup.py

@@ -51,7 +51,7 @@ install_requires = [
     "configobj>=5.0.6",
     "gitpython>3",
     "dulwich>=0.20.20",
-    "pygit2>=1.4.0",
+    "pygit2>=1.5.0",
     "setuptools>=34.0.0",
     "nanotime>=0.5.2",
     "pyasn1>=0.4.1",

tests/func/test_scm.py

@@ -1,6 +1,7 @@
 import os
 
 import pytest
+from flaky.flaky_decorator import flaky
 from git import Repo
 
 from dvc.scm import SCM, Git, NoSCM
@@ -242,6 +243,9 @@ def test_git_stash_drop(tmp_dir, scm, ref):
     assert len(stash) == 1
 
 
+# libgit2 stash_save() is flaky on linux when run inside pytest, see:
+# https://github.com/iterative/dvc/pull/5286#issuecomment-792574294
+@flaky(max_runs=5, min_passes=1)
 @pytest.mark.parametrize("ref", [None, "refs/foo/stash"])
 def test_git_stash_pop(tmp_dir, scm, ref):
     from dvc.scm.git import Stash
@@ -279,6 +283,7 @@ def test_git_stash_clear(tmp_dir, scm, ref):
 
     parts = list(stash.ref.split("/"))
     assert not os.path.exists(os.path.join(os.fspath(tmp_dir), ".git", *parts))
-    assert not os.path.exists(
-        os.path.join(os.fspath(tmp_dir), ".git", "logs", *parts)
-    )
+    # NOTE: some backends will completely remove reflog file on clear, some
+    # will only truncate it, either case means an empty stash
+    log_path = os.path.join(os.fspath(tmp_dir), ".git", "logs", *parts)
+    assert not os.path.exists(log_path) or not open(log_path).read()

tests/unit/scm/test_git.py

@@ -165,6 +165,25 @@ def test_set_ref(tmp_dir, git):
     )
 
 
+def test_set_ref_with_message(tmp_dir, git):
+    tmp_dir.scm_gen({"file": "0"}, commit="init")
+    init_rev = tmp_dir.scm.get_rev()
+    tmp_dir.scm_gen({"file": "1"}, commit="commit")
+    commit_rev = tmp_dir.scm.get_rev()
+
+    git.set_ref("refs/foo/bar", init_rev, message="init message")
+    assert (
+        "init message"
+        in (tmp_dir / ".git" / "logs" / "refs" / "foo" / "bar").read_text()
+    )
+
+    git.set_ref("refs/foo/bar", commit_rev, message="modify message")
+    assert (
+        "modify message"
+        in (tmp_dir / ".git" / "logs" / "refs" / "foo" / "bar").read_text()
+    )
+
+
 def test_get_ref(tmp_dir, git):
     tmp_dir.scm_gen({"file": "0"}, commit="init")
     init_rev = tmp_dir.scm.get_rev()