From 90a636e33f79f58e466083d5fc1c971ee46cfb95 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <ruslan@iterative.ai>
Date: Sun, 21 Oct 2018 00:19:33 +0300
Subject: [PATCH] scm: allow all_tags and all_branches

Fixes #1243

Signed-off-by: Ruslan Kuprieiev <ruslan@iterative.ai>
---
 dvc/scm.py       | 27 ++++++++++++++--------
 tests/test_gc.py | 60 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/dvc/scm.py b/dvc/scm.py
index 2c7f35b1d..742feb4f7 100644
--- a/dvc/scm.py
+++ b/dvc/scm.py
@@ -50,6 +50,9 @@ class Base(object):
     def branch(self, branch):
         pass
 
+    def tag(self, tag):
+        pass
+
     def brancher(self,
                  branches=None,
                  all_branches=False,
@@ -63,16 +66,19 @@ class Base(object):
         saved = self.active_branch()
         revs = []
 
-        if branches is not None:
-            revs.extend(branches)
-        elif all_branches:
-            revs.extend(self.list_branches())
-        elif tags is not None:
-            revs.extend(tags)
-        elif all_tags:
-            revs.extend(self.list_tags())
-        else:
+        if all_branches:
+            branches = self.list_branches()
+
+        if all_tags:
+            tags = self.list_tags()
+
+        if branches is None:
             revs.extend([saved])
+        else:
+            revs.extend(branches)
+
+        if tags is not None:
+            revs.extend(tags)
 
         for rev in revs:
             self.checkout(rev)
@@ -215,6 +221,9 @@ class Git(Base):
     def branch(self, branch):
         self.repo.git.branch(branch)
 
+    def tag(self, tag):
+        self.repo.git.tag(tag)
+
     def untracked_files(self):
         files = self.repo.untracked_files
         return [os.path.join(self.repo.working_dir, fname) for fname in files]
diff --git a/tests/test_gc.py b/tests/test_gc.py
index 52da378be..dcafd6005 100644
--- a/tests/test_gc.py
+++ b/tests/test_gc.py
@@ -40,3 +40,63 @@ class TestGC(TestDvc):
 
         for c in self.good_cache:
             self.assertTrue(os.path.exists(c))
+
+
+class TestGCBranchesTags(TestDvc):
+    def _check_cache(self, num):
+        total = 0
+        for root, dirs, files in os.walk(os.path.join('.dvc', 'cache')):
+            total += len(files)
+        self.assertEqual(total, num)
+
+    def test(self):
+        fname = 'file'
+
+        with open(fname, 'w+') as fobj:
+            fobj.write('v1.0')
+
+        stages = self.dvc.add(fname)
+        self.assertEqual(len(stages), 1)
+        self.dvc.scm.add(['.gitignore', stages[0].relpath])
+        self.dvc.scm.commit('v1.0')
+        self.dvc.scm.tag('v1.0')
+
+        self.dvc.scm.checkout('test', create_new=True)
+        self.dvc.remove(stages[0].relpath, outs_only=True)
+        with open(fname, 'w+') as fobj:
+            fobj.write('test')
+        stages = self.dvc.add(fname)
+        self.assertEqual(len(stages), 1)
+        self.dvc.scm.add(['.gitignore', stages[0].relpath])
+        self.dvc.scm.commit('test')
+
+        self.dvc.scm.checkout('master')
+        self.dvc.remove(stages[0].relpath, outs_only=True)
+        with open(fname, 'w+') as fobj:
+            fobj.write('trash')
+        stages = self.dvc.add(fname)
+        self.assertEqual(len(stages), 1)
+        self.dvc.scm.add(['.gitignore', stages[0].relpath])
+        self.dvc.scm.commit('trash')
+
+        self.dvc.remove(stages[0].relpath, outs_only=True)
+        with open(fname, 'w+') as fobj:
+            fobj.write('master')
+        stages = self.dvc.add(fname)
+        self.assertEqual(len(stages), 1)
+        self.dvc.scm.add(['.gitignore', stages[0].relpath])
+        self.dvc.scm.commit('master')
+
+        self._check_cache(4)
+
+        self.dvc.gc(all_tags=True, all_branches=True)
+
+        self._check_cache(3)
+
+        self.dvc.gc(all_tags=False, all_branches=True)
+
+        self._check_cache(2)
+
+        self.dvc.gc(all_tags=True, all_branches=False)
+
+        self._check_cache(1)