add: cleanup .gitignore decorator

dvc/repo/__init__.py

@@ -62,8 +62,6 @@ def __init__(self, root_dir=None):
 
         self.metrics = Metrics(self)
 
-        self.files_to_git_add = []
-
         self._ignore()
 
         self.updater.check()
@@ -92,17 +90,6 @@ def find_dvc_dir(root=None):
         root_dir = Repo.find_root(root)
         return os.path.join(root_dir, Repo.DVC_DIR)
 
-    def remind_to_git_add(self):
-        if not self.files_to_git_add:
-            return
-
-        logger.info(
-            "\n"
-            "To track the changes with git run:\n"
-            "\n"
-            "\tgit add {files}".format(files=" ".join(self.files_to_git_add))
-        )
-
     @staticmethod
     def init(root_dir=os.curdir, no_scm=False, force=False):
         from dvc.repo.init import init

dvc/repo/add.py

@@ -1,9 +1,12 @@
 import os
+
+from dvc.scm import scm_context
 from dvc.stage import Stage
 from dvc.utils import walk_files
 from dvc.exceptions import RecursiveAddingWhileUsingFilename
 
 
+@scm_context
 def add(repo, target, recursive=False, no_commit=False, fname=None):
     if recursive and fname:
         raise RecursiveAddingWhileUsingFilename()
@@ -25,9 +28,6 @@ def add(repo, target, recursive=False, no_commit=False, fname=None):
 
     for stage in stages:
         stage.dump()
-
-    repo.remind_to_git_add()
-
     return stages
 
 

dvc/repo/imp.py

@@ -1,3 +1,7 @@
+from dvc.scm import scm_context
+
+
+@scm_context
 def imp(self, url, out, resume=False):
     from dvc.stage import Stage
 
@@ -9,13 +13,9 @@ def imp(self, url, out, resume=False):
 
     self.check_dag(self.stages() + [stage])
 
-    self.files_to_git_add = []
-
     with self.state:
         stage.run(resume=resume)
 
     stage.dump()
 
-    self.remind_to_git_add()
-
     return stage

dvc/repo/move.py

@@ -1,6 +1,7 @@
 import os
 
 from dvc.exceptions import MoveNotDataSourceError
+from dvc.scm import scm_context
 
 
 def _expand_target_path(from_path, to_path):
@@ -9,6 +10,7 @@ def _expand_target_path(from_path, to_path):
     return to_path
 
 
+@scm_context
 def move(self, from_path, to_path):
     """
     Renames an output file and modifies the stage associated
@@ -64,5 +66,3 @@ def move(self, from_path, to_path):
         out.move(to_out)
 
     stage.dump()
-
-    self.remind_to_git_add()

dvc/repo/reproduce.py

@@ -4,6 +4,7 @@
 
 import dvc.logger as logger
 from dvc.exceptions import ReproductionError
+from dvc.scm import scm_context
 
 
 def _reproduce_stage(stages, node, force, dry, interactive, no_commit):
@@ -27,6 +28,7 @@ def _reproduce_stage(stages, node, force, dry, interactive, no_commit):
     return [stage]
 
 
+@scm_context
 def reproduce(
     self,
     target=None,
@@ -65,8 +67,6 @@ def reproduce(
     else:
         targets.append(target)
 
-    self.files_to_git_add = []
-
     ret = []
     with self.state:
         for target in targets:
@@ -82,8 +82,6 @@ def reproduce(
             )
             ret.extend(stages)
 
-    self.remind_to_git_add()
-
     return ret
 
 

dvc/repo/run.py

@@ -1,6 +1,9 @@
 from __future__ import unicode_literals
 
+from dvc.scm import scm_context
 
+
+@scm_context
 def run(
     self,
     cmd=None,
@@ -53,13 +56,10 @@ def run(
 
     self.check_dag(self.stages() + [stage])
 
-    self.files_to_git_add = []
     with self.state:
         if not no_exec:
             stage.run(no_commit=no_commit)
 
     stage.dump()
 
-    self.remind_to_git_add()
-
     return stage

dvc/scm/__init__.py

@@ -5,6 +5,8 @@
 from dvc.scm.base import Base
 from dvc.scm.git import Git
 
+import dvc
+
 
 def SCM(root_dir, repo=None):  # pylint: disable=invalid-name
     """Returns SCM instance that corresponds to a repo at the specified
@@ -21,3 +23,19 @@ def SCM(root_dir, repo=None):  # pylint: disable=invalid-name
         return Git(root_dir, repo=repo)
 
     return Base(root_dir, repo=repo)
+
+
+def scm_context(method):
+    def run(*args, **kw):
+        repo = args[0]
+        assert isinstance(repo, dvc.repo.Repo)
+        try:
+            result = method(*args, **kw)
+            repo.scm.reset_ignores()
+            repo.scm.remind_to_track()
+            return result
+        except Exception as e:
+            repo.scm.cleanup_ignores()
+            raise e
+
+    return run

dvc/scm/base.py

@@ -139,3 +139,26 @@ def list_tags(self):  # pylint: disable=no-self-use
 
     def install(self):
         """Adds dvc commands to SCM hooks for the repo."""
+
+    def cleanup_ignores(self):
+        """
+        This method should clean up ignores (eg. entries in .gitignore),
+        use, when method editing ignores (eg. add, run, import) fails to
+        perform its task.
+        """
+
+    def reset_ignores(self):
+        """
+        Method to reset in-memory ignore storing mechanism.
+        """
+
+    def remind_to_track(self):
+        """
+        Method to remind user to track newly created files handled by scm
+        """
+
+    def track_file(self, path):
+        """
+        Method to add file to mechanism that will remind user
+        to track new files
+        """

dvc/scm/git.py

@@ -39,6 +39,9 @@ def __init__(self, root_dir=os.curdir, repo=None):
         libpath = env.get("LD_LIBRARY_PATH", None)
         self.git.git.update_environment(LD_LIBRARY_PATH=libpath)
 
+        self.ignored_paths = []
+        self.files_to_track = []
+
     @staticmethod
     def is_repo(root_dir):
         return os.path.isdir(Git._get_git_dir(root_dir))
@@ -95,8 +98,9 @@ def ignore(self, path):
         with open(gitignore, "a") as fobj:
             fobj.write(content)
 
-        if self.repo is not None:
-            self.repo.files_to_git_add.append(os.path.relpath(gitignore))
+        self.track_file(os.path.relpath(gitignore))
+
+        self.ignored_paths.append(path)
 
     def ignore_remove(self, path):
         entry, gitignore = self._get_gitignore(path)
@@ -112,8 +116,7 @@ def ignore_remove(self, path):
         with open(gitignore, "w") as fobj:
             fobj.writelines(filtered)
 
-        if self.repo is not None:
-            self.repo.files_to_git_add.append(os.path.relpath(gitignore))
+        self.track_file(os.path.relpath(gitignore))
 
     def add(self, paths):
         # NOTE: GitPython is not currently able to handle index version >= 3.
@@ -173,3 +176,25 @@ def _install_hook(self, name, cmd):
     def install(self):
         self._install_hook("post-checkout", "checkout")
         self._install_hook("pre-commit", "status")
+
+    def cleanup_ignores(self):
+        for path in self.ignored_paths:
+            self.ignore_remove(path)
+        self.reset_ignores()
+
+    def reset_ignores(self):
+        self.ignored_paths = []
+
+    def remind_to_track(self):
+        if not self.files_to_track:
+            return
+
+        logger.info(
+            "\n"
+            "To track the changes with git run:\n"
+            "\n"
+            "\tgit add {files}".format(files=" ".join(self.files_to_track))
+        )
+
+    def track_file(self, path):
+        self.files_to_track.append(path)

dvc/stage.py

@@ -552,7 +552,7 @@ def dump(self):
         with open(fname, "w") as fd:
             yaml.safe_dump(d, fd, default_flow_style=False)
 
-        self.repo.files_to_git_add.append(os.path.relpath(fname))
+        self.repo.scm.track_file(os.path.relpath(fname))
 
     def _compute_md5(self):
         from dvc.output.local import OutputLOCAL

tests/test_add.py

@@ -21,7 +21,7 @@
 from dvc.repo import Repo as DvcRepo
 
 from tests.basic_env import TestDvc
-from tests.utils import spy
+from tests.utils import spy, get_gitignore_content
 from tests.utils.logger import MockLoggerHandlers, ConsoleFontColorsRemover
 
 
@@ -457,3 +457,23 @@ def test(self):
         self.assertEqual(0, ret)
         self.assertTrue(os.path.exists("bar.dvc"))
         self.assertFalse(os.path.exists("foo.dvc"))
+
+
+class TestShouldCleanUpAfterFailedAdd(TestDvc):
+    def test(self):
+        ret = main(["add", self.FOO])
+        self.assertEqual(0, ret)
+
+        foo_stage_file = self.FOO + Stage.STAGE_FILE_SUFFIX
+        # corrupt stage file
+        with open(foo_stage_file, "a+") as file:
+            file.write("this will break yaml file structure")
+
+        ret = main(["add", self.BAR])
+        self.assertEqual(1, ret)
+
+        bar_stage_file = self.BAR + Stage.STAGE_FILE_SUFFIX
+        self.assertFalse(os.path.exists(bar_stage_file))
+
+        gitignore_content = get_gitignore_content()
+        self.assertFalse(any(self.BAR in line for line in gitignore_content))

tests/test_scm.py

@@ -5,6 +5,7 @@
 from dvc.scm import SCM, Base, Git
 
 from tests.basic_env import TestDir, TestGit, TestGitSubmodule
+from tests.utils import get_gitignore_content
 
 
 class TestSCM(TestDir):
@@ -43,11 +44,9 @@ def test_commit_in_submodule(self):
 
 class TestIgnore(TestGit):
     def _count_gitignore(self):
-        with open(Git.GITIGNORE, "r") as fd:
-            lines = fd.readlines()
-            return len(
-                list(filter(lambda x: x.strip() == "/" + self.FOO, lines))
-            )
+        lines = get_gitignore_content()
+
+        return len(list(filter(lambda x: x.strip() == "/" + self.FOO, lines)))
 
     def test(self):
         git = Git(self._root_dir)

tests/unit/scm/scm.py

@@ -0,0 +1,36 @@
+from dvc.repo import Repo
+from dvc.scm import scm_context, Base
+from unittest import TestCase
+import mock
+
+
+class TestScmContext(TestCase):
+    def setUp(self):
+        self.repo_mock = mock.Mock(spec=Repo)
+        self.scm_mock = mock.Mock(spec=Base)
+        self.repo_mock.scm = self.scm_mock
+
+    def test_should_successfully_perform_method(self):
+        method = mock.Mock()
+        wrapped = scm_context(method)
+
+        wrapped(self.repo_mock)
+
+        self.assertEqual(1, method.call_count)
+        self.assertEqual(1, self.scm_mock.reset_ignores.call_count)
+        self.assertEqual(1, self.scm_mock.remind_to_track.call_count)
+
+        self.assertEqual(0, self.scm_mock.cleanup_ignores.call_count)
+
+    def test_should_throw_and_cleanup(self):
+        method = mock.Mock(side_effect=Exception("some problem"))
+        wrapped = scm_context(method)
+
+        with self.assertRaises(Exception):
+            wrapped(self.repo_mock)
+
+        self.assertEqual(1, method.call_count)
+        self.assertEqual(1, self.scm_mock.cleanup_ignores.call_count)
+
+        self.assertEqual(0, self.scm_mock.reset_ignores.call_count)
+        self.assertEqual(0, self.scm_mock.remind_to_track.call_count)

tests/utils/__init__.py

@@ -1,3 +1,4 @@
+import yaml
 import os
 
 from dvc.scm import Git
@@ -21,6 +22,11 @@ def get_gitignore_content():
         return gitignore.readlines()
 
 
+def load_stage_file(path):
+    with open(path, "r") as fobj:
+        return yaml.safe_load(fobj)
+
+
 @contextmanager
 def cd(newdir):
     prevdir = os.getcwd()