show git-add hint on stage:add, refactor scm_context (#5290)

* Refactor scm_context, show git-add hint on stage:add

* move tests

dvc/command/stage.py

@@ -107,11 +107,13 @@ class CmdStageList(CmdBase):
 
 class CmdStageAdd(CmdBase):
     def run(self):
+        repo = self.repo
         kwargs = vars(self.args)
-        stage = self.repo.stage.create_from_cli(validate=True, **kwargs)
+        stage = repo.stage.create_from_cli(validate=True, **kwargs)
 
-        stage.ignore_outs()
-        stage.dump()
+        with repo.scm.track_file_changes(config=repo.config):
+            stage.ignore_outs()
+            stage.dump()
         return 0
 
 

dvc/repo/scm_context.py

@@ -1,17 +1,16 @@
+from functools import wraps
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from dvc.repo import Repo
+
+
 def scm_context(method):
-    def run(repo, *args, **kw):
-        try:
-            result = method(repo, *args, **kw)
-            repo.scm.reset_ignores()
-            autostage = repo.config.get("core", {}).get("autostage", False)
-            if autostage:
-                repo.scm.track_changed_files()
-            else:
-                repo.scm.remind_to_track()
-            repo.scm.reset_tracked_files()
-            return result
-        except Exception:
-            repo.scm.cleanup_ignores()
-            raise
+    @wraps(method)
+    def run(repo: "Repo", *args, **kw):
+        scm = repo.scm
+
+        with scm.track_file_changes(config=repo.config):
+            return method(repo, *args, **kw)
 
     return run

dvc/scm/base.py

@@ -1,6 +1,9 @@
 """Manages source control systems (e.g. Git) in DVC."""
 
 import os
+from contextlib import contextmanager
+
+from funcy import get_in
 
 from dvc.exceptions import DvcException
 
@@ -173,6 +176,23 @@ class Base:
         to track new files
         """
 
+    @contextmanager
+    def track_file_changes(self, config=None):
+        autostage = get_in(config or {}, ["core", "autostage"])
+
+        try:
+            yield
+        except Exception:
+            self.cleanup_ignores()
+            raise
+
+        self.reset_ignores()
+        if autostage:
+            self.track_changed_files()
+        else:
+            self.remind_to_track()
+        self.reset_tracked_files()
+
     def belongs_to_scm(self, path):
         """Return boolean whether file belongs to scm"""
 

tests/unit/scm/test_git.py

@@ -498,3 +498,9 @@ def test_reset(tmp_dir, scm, git):
     staged, unstaged, _ = scm.status()
     assert len(staged) == 1
     assert len(unstaged) == 1
+
+
+def test_remind_to_track(scm, caplog):
+    scm.files_to_track = ["fname with spaces.txt", "тест", "foo"]
+    scm.remind_to_track()
+    assert "git add 'fname with spaces.txt' 'тест' foo" in caplog.text

tests/unit/scm/test_scm.py

@@ -1,60 +1,61 @@
-from unittest import TestCase
+import pytest
 
-import mock
-
-from dvc.repo import Repo
 from dvc.repo.scm_context import scm_context
 from dvc.scm import NoSCM
 
 
-class TestScmContext(TestCase):
-    def setUp(self):
-        self.repo_mock = mock.Mock(spec=Repo)
-        self.scm_mock = mock.Mock(spec=NoSCM)
-        self.repo_mock.scm = self.scm_mock
+def test_scm_context(dvc, mocker):
+    scm = dvc.scm
+    m = mocker.patch.object(
+        scm, "track_file_changes", wraps=scm.track_file_changes
+    )
 
-    def test_should_successfully_perform_method(self):
-        method = mock.Mock()
-        wrapped = scm_context(method)
+    ret_val = 0
+    wrapped = scm_context(mocker.Mock(return_value=ret_val))
 
-        self.repo_mock.configure_mock(config={})
-        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_check_autostage(self):
-        method = mock.Mock()
-        wrapped = scm_context(method)
-
-        config_autostage_attrs = {"config": {"core": {"autostage": True}}}
-        self.repo_mock.configure_mock(**config_autostage_attrs)
-        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.track_changed_files.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)
+    assert wrapped(dvc) == ret_val
+    m.assert_called_once_with(config=dvc.config)
 
 
-def test_remind_to_track(scm, caplog):
-    scm.files_to_track = ["fname with spaces.txt", "тест", "foo"]
-    scm.remind_to_track()
-    assert "git add 'fname with spaces.txt' 'тест' foo" in caplog.text
+def test_track_file_changes(mocker):
+    scm = mocker.Mock(NoSCM)
+
+    with NoSCM.track_file_changes(scm):
+        pass
+
+    assert scm.reset_ignores.call_count == 1
+    assert scm.remind_to_track.call_count == 1
+    assert scm.track_changed_files.call_count == 0
+    assert scm.cleanup_ignores.call_count == 0
+    assert scm.reset_tracked_files.call_count == 1
+
+
+def test_track_file_changes_autostage(mocker):
+    scm = mocker.Mock(NoSCM)
+
+    config = {"core": {"autostage": True}}
+    with NoSCM.track_file_changes(scm, config=config):
+        pass
+
+    assert scm.track_changed_files.call_count == 1
+    assert scm.reset_ignores.call_count == 1
+    assert scm.remind_to_track.call_count == 0
+    assert scm.cleanup_ignores.call_count == 0
+    assert scm.reset_tracked_files.call_count == 1
+
+
+def test_track_file_changes_throw_and_cleanup(mocker):
+    scm = mocker.Mock(NoSCM)
+
+    class CustomException(Exception):
+        pass
+
+    with pytest.raises(CustomException, match="oops"):
+        with NoSCM.track_file_changes(scm):
+            raise CustomException("oops")
+
+    assert scm.cleanup_ignores.call_count == 1
+    assert scm.reset_ignores.call_count == 0
+    assert scm.track_changed_files.call_count == 0
+    assert scm.remind_to_track.call_count == 0
+    assert scm.reset_tracked_files.call_count == 0