From 98590779eb2272a819318485a47bd99c9d081340 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?=
 <peter@pmrowla.com>
Date: Fri, 19 Feb 2021 16:42:04 +0900
Subject: [PATCH] pygit: implement get_refs_containing (#5495)

* git: make is_sha a Git class method instead of backend function

* pygit: implement get_refs_containing

* use merge base instead of parents
---
 dvc/scm/git/__init__.py        |  9 +++----
 dvc/scm/git/backend/base.py    |  5 ----
 dvc/scm/git/backend/dulwich.py |  4 ----
 dvc/scm/git/backend/pygit2.py  | 43 +++++++++++++++++++++++++++-------
 4 files changed, 38 insertions(+), 23 deletions(-)

diff --git a/dvc/scm/git/__init__.py b/dvc/scm/git/__init__.py
index 6d4868f3d..65e05d2b3 100644
--- a/dvc/scm/git/__init__.py
+++ b/dvc/scm/git/__init__.py
@@ -2,6 +2,7 @@
 
 import logging
 import os
+import re
 import shlex
 from collections.abc import Mapping
 from contextlib import contextmanager
@@ -75,6 +76,7 @@ class Git(Base):
     GITIGNORE = ".gitignore"
     GIT_DIR = ".git"
     LOCAL_BRANCH_PREFIX = "refs/heads/"
+    RE_HEXSHA = re.compile(r"^[0-9A-Fa-f]{4,40}$")
 
     def __init__(
         self, *args, backends: Optional[Iterable[str]] = None, **kwargs
@@ -119,12 +121,7 @@ class Git(Base):
 
     @classmethod
     def is_sha(cls, rev):
-        for _, backend in GitBackends.DEFAULT.items():
-            try:
-                return backend.is_sha(rev)
-            except NotImplementedError:
-                pass
-        raise NoGitBackendError("is_sha")
+        return rev and cls.RE_HEXSHA.search(rev)
 
     @staticmethod
     def _get_git_dir(root_dir):
diff --git a/dvc/scm/git/backend/base.py b/dvc/scm/git/backend/base.py
index 6fe9af08e..63ac4457e 100644
--- a/dvc/scm/git/backend/base.py
+++ b/dvc/scm/git/backend/base.py
@@ -41,11 +41,6 @@ class BaseGitBackend(ABC):
     ):
         pass
 
-    @staticmethod
-    @abstractmethod
-    def is_sha(rev: str) -> bool:
-        pass
-
     @property
     @abstractmethod
     def dir(self) -> str:
diff --git a/dvc/scm/git/backend/dulwich.py b/dvc/scm/git/backend/dulwich.py
index b628eac82..9f496ab45 100644
--- a/dvc/scm/git/backend/dulwich.py
+++ b/dvc/scm/git/backend/dulwich.py
@@ -118,10 +118,6 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
     ):
         raise NotImplementedError
 
-    @staticmethod
-    def is_sha(rev: str) -> bool:
-        raise NotImplementedError
-
     @property
     def dir(self) -> str:
         return self.repo.commondir()
diff --git a/dvc/scm/git/backend/pygit2.py b/dvc/scm/git/backend/pygit2.py
index 45af40173..d7c816ed8 100644
--- a/dvc/scm/git/backend/pygit2.py
+++ b/dvc/scm/git/backend/pygit2.py
@@ -76,10 +76,6 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
     ):
         raise NotImplementedError
 
-    @staticmethod
-    def is_sha(rev: str) -> bool:
-        raise NotImplementedError
-
     @property
     def dir(self) -> str:
         raise NotImplementedError
@@ -225,12 +221,43 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
         ref.delete()
 
     def iter_refs(self, base: Optional[str] = None):
-        for ref in self.repo.references:
-            if ref.startswith(base):
-                yield ref
+        if base:
+            for ref in self.repo.references:
+                if ref.startswith(base):
+                    yield ref
+        else:
+            yield from self.repo.references
 
     def get_refs_containing(self, rev: str, pattern: Optional[str] = None):
-        raise NotImplementedError
+        import fnmatch
+
+        from pygit2 import GitError
+
+        def _contains(repo, ref, search_commit):
+            commit, _ref = self.repo.resolve_refish(ref)
+            base = repo.merge_base(search_commit.id, commit.id)
+            return base == search_commit.id
+
+        try:
+            search_commit, _ref = self.repo.resolve_refish(rev)
+        except (KeyError, GitError):
+            raise SCMError(f"Invalid rev '{rev}'")
+
+        if not pattern:
+            yield from (
+                ref
+                for ref in self.iter_refs()
+                if _contains(self.repo, ref, search_commit)
+            )
+            return
+
+        literal = pattern.rstrip("/").split("/")
+        for ref in self.iter_refs():
+            if (
+                ref.split("/")[: len(literal)] == literal
+                or fnmatch.fnmatch(ref, pattern)
+            ) and _contains(self.repo, ref, search_commit):
+                yield ref
 
     def push_refspec(
         self,