From 3d58f4b8ce234ca1d3762865b12677afef4e4e63 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, 30 Oct 2020 19:06:35 +0900
Subject: [PATCH] git: move branch iteration functionality from experiments
 into scm.git (#4810)

* git: move branch iteration functionality from experiments into scm.git

* add branch_revs unit test
---
 dvc/repo/experiments/show.py | 30 ++----------------------------
 dvc/scm/git.py               | 25 ++++++++++++++++++++++++-
 tests/unit/scm/test_git.py   | 14 ++++++++++++++
 3 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/dvc/repo/experiments/show.py b/dvc/repo/experiments/show.py
index a8fa630c7..688fc6687 100644
--- a/dvc/repo/experiments/show.py
+++ b/dvc/repo/experiments/show.py
@@ -1,9 +1,6 @@
 import logging
 from collections import OrderedDict, defaultdict
 from datetime import datetime
-from typing import Optional
-
-from funcy import first
 
 from dvc.repo import locked
 from dvc.repo.metrics.show import _collect_metrics, _read_metrics
@@ -20,7 +17,7 @@ def _collect_experiment(repo, rev, stash=False, sha_only=True):
         if rev == "workspace":
             res["timestamp"] = None
         else:
-            commit = _resolve_commit(repo, rev)
+            commit = repo.scm.resolve_commit(rev)
             res["timestamp"] = datetime.fromtimestamp(commit.committed_date)
 
         configs = _collect_configs(repo, rev=rev)
@@ -47,38 +44,15 @@ def _collect_experiment(repo, rev, stash=False, sha_only=True):
     return res
 
 
-def _resolve_commit(repo, rev):
-    from git.objects.tag import TagObject
-
-    commit = repo.scm.repo.rev_parse(rev)
-    if isinstance(commit, TagObject):
-        commit = commit.object
-    return commit
-
-
 def _collect_checkpoint_experiment(repo, branch, baseline, **kwargs):
     res = OrderedDict()
     exp_rev = repo.scm.resolve_rev(branch)
-    for rev in _branch_revs(repo, exp_rev, baseline):
+    for rev in repo.scm.branch_revs(exp_rev, baseline):
         res[rev] = _collect_experiment(repo, rev, **kwargs)
         res[rev]["checkpoint_tip"] = exp_rev
     return res
 
 
-def _branch_revs(repo, branch_tip, baseline: Optional[str] = None):
-    """Iterate over revisions in a given branch (from newest to oldest).
-
-    If baseline is set, iterator will stop when the specified revision is
-    reached.
-    """
-    commit = _resolve_commit(repo, branch_tip)
-    while commit is not None:
-        yield commit.hexsha
-        commit = first(commit.parents)
-        if commit and commit.hexsha == baseline:
-            return
-
-
 @locked
 def show(
     repo,
diff --git a/dvc/scm/git.py b/dvc/scm/git.py
index 5b3965f15..5bfd84c77 100644
--- a/dvc/scm/git.py
+++ b/dvc/scm/git.py
@@ -4,8 +4,9 @@ import logging
 import os
 import shlex
 from functools import partial
+from typing import Optional
 
-from funcy import cached_property
+from funcy import cached_property, first
 from pathspec.patterns import GitWildMatchPattern
 
 from dvc.exceptions import GitHookAlreadyExistsError
@@ -468,3 +469,25 @@ class Git(Base):
     @property
     def no_commits(self):
         return not self.list_all_commits()
+
+    def branch_revs(self, branch: str, end_rev: Optional[str] = None):
+        """Iterate over revisions in a given branch (from newest to oldest).
+
+        If end_rev is set, iterator will stop when the specified revision is
+        reached.
+        """
+        commit = self.resolve_commit(branch)
+        while commit is not None:
+            yield commit.hexsha
+            commit = first(commit.parents)
+            if commit and commit.hexsha == end_rev:
+                return
+
+    def resolve_commit(self, rev):
+        """Return Commit object for the specified revision."""
+        from git.objects.tag import TagObject
+
+        commit = self.repo.rev_parse(rev)
+        if isinstance(commit, TagObject):
+            commit = commit.object
+        return commit
diff --git a/tests/unit/scm/test_git.py b/tests/unit/scm/test_git.py
index 2f82dd04e..0c30db04f 100644
--- a/tests/unit/scm/test_git.py
+++ b/tests/unit/scm/test_git.py
@@ -109,3 +109,17 @@ def test_no_commits(tmp_dir):
     Git().commit("foo")
 
     assert not Git().no_commits
+
+
+def test_branch_revs(tmp_dir, scm):
+    tmp_dir.scm_gen({"file": "0"}, commit="init")
+    init_rev = scm.get_rev()
+
+    expected = []
+    for i in range(1, 5):
+        tmp_dir.scm_gen({"file": f"{i}"}, commit=f"{i}")
+        expected.append(scm.get_rev())
+
+    for rev in scm.branch_revs("master", init_rev):
+        assert rev == expected.pop()
+    assert len(expected) == 0