Repository: iterative/dvc

Issue:
run: handle SIGINT gracefully
https://discordapp.com/channels/485586884165107732/485596304961962003/543461241289572352

Discussion:

Patch:
diff --git a/dvc/repo/brancher.py b/dvc/repo/brancher.py
index 72b066bee..e84fcaf70 100644
--- a/dvc/repo/brancher.py
+++ b/dvc/repo/brancher.py
@@ -1,6 +1,3 @@
-from dvc.scm.tree import WorkingTree
-
-
 def brancher(  # noqa: E302
     self, branches=None, all_branches=False, tags=None, all_tags=False
 ):
@@ -27,8 +24,11 @@ def brancher(  # noqa: E302
 
     scm = self.scm
 
-    self.tree = WorkingTree(self.root_dir)
-    yield "working tree"
+    if self.scm.is_dirty():
+        from dvc.scm.tree import WorkingTree
+
+        self.tree = WorkingTree(self.root_dir)
+        yield "Working Tree"
 
     if all_branches:
         branches = scm.list_branches()
@@ -36,7 +36,9 @@ def brancher(  # noqa: E302
     if all_tags:
         tags = scm.list_tags()
 
-    if branches is not None:
+    if branches is None:
+        revs.extend([scm.active_branch()])
+    else:
         revs.extend(branches)
 
     if tags is not None:
diff --git a/dvc/stage.py b/dvc/stage.py
index 93831657e..7610653f3 100644
--- a/dvc/stage.py
+++ b/dvc/stage.py
@@ -7,6 +7,7 @@ import re
 import os
 import subprocess
 import logging
+import signal
 
 from dvc.utils import relpath
 from dvc.utils.compat import pathlib
@@ -767,16 +768,22 @@ class Stage(object):
         executable = os.getenv("SHELL") if os.name != "nt" else None
         self._warn_if_fish(executable)
 
-        p = subprocess.Popen(
-            self.cmd,
-            cwd=self.wdir,
-            shell=True,
-            env=fix_env(os.environ),
-            executable=executable,
-        )
-        p.communicate()
+        old_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
+        p = None
+
+        try:
+            p = subprocess.Popen(
+                self.cmd,
+                cwd=self.wdir,
+                shell=True,
+                env=fix_env(os.environ),
+                executable=executable,
+            )
+            p.communicate()
+        finally:
+            signal.signal(signal.SIGINT, old_handler)
 
-        if p.returncode != 0:
+        if (p is None) or (p.returncode != 0):
             raise StageCmdFailedError(self)
 
     def run(self, dry=False, resume=False, no_commit=False, force=False):


Test patch:
diff --git a/tests/func/test_metrics.py b/tests/func/test_metrics.py
index ab137fd99..806afcca7 100644
--- a/tests/func/test_metrics.py
+++ b/tests/func/test_metrics.py
@@ -706,7 +706,6 @@ class TestCachedMetrics(TestDvc):
                 "master": {"metrics.json": ["master"]},
                 "one": {"metrics.json": ["one"]},
                 "two": {"metrics.json": ["two"]},
-                "working tree": {"metrics.json": ["two"]},
             },
         )
 
@@ -720,7 +719,6 @@ class TestCachedMetrics(TestDvc):
                 "master": {"metrics.json": ["master"]},
                 "one": {"metrics.json": ["one"]},
                 "two": {"metrics.json": ["two"]},
-                "working tree": {"metrics.json": ["two"]},
             },
         )
 
diff --git a/tests/func/test_run.py b/tests/func/test_run.py
index 79ae192db..3d707cdda 100644
--- a/tests/func/test_run.py
+++ b/tests/func/test_run.py
@@ -6,6 +6,7 @@ import mock
 import shutil
 import filecmp
 import subprocess
+import signal
 
 from dvc.main import main
 from dvc.output import OutputBase
@@ -323,9 +324,46 @@ class TestCmdRun(TestDvc):
         self.assertEqual(ret, 0)
         self.assertEqual(stage.cmd, 'echo "foo bar"')
 
-    @mock.patch.object(subprocess, "Popen", side_effect=KeyboardInterrupt)
-    def test_keyboard_interrupt(self, _):
-        ret = main(["run", "mycmd"])
+    @mock.patch.object(subprocess.Popen, "wait", new=KeyboardInterrupt)
+    def test_keyboard_interrupt(self):
+        ret = main(
+            [
+                "run",
+                "-d",
+                self.FOO,
+                "-d",
+                self.CODE,
+                "-o",
+                "out",
+                "-f",
+                "out.dvc",
+                "python",
+                self.CODE,
+                self.FOO,
+                "out",
+            ]
+        )
+        self.assertEqual(ret, 1)
+
+    @mock.patch.object(signal, "signal", side_effect=[None, KeyboardInterrupt])
+    def test_keyboard_interrupt_after_second_signal_call(self, _):
+        ret = main(
+            [
+                "run",
+                "-d",
+                self.FOO,
+                "-d",
+                self.CODE,
+                "-o",
+                "out",
+                "-f",
+                "out.dvc",
+                "python",
+                self.CODE,
+                self.FOO,
+                "out",
+            ]
+        )
         self.assertEqual(ret, 252)