added old signal persistance and fixed tests accordingly

dvc/stage.py

@@ -768,19 +768,22 @@ def _run(self):
         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,
-        )
+        old_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
+        p = None
 
-        signal.signal(signal.SIGINT, signal.SIG_IGN)
-        p.communicate()
-        signal.signal(signal.SIGINT, signal.default_int_handler)
+        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):

tests/func/test_run.py

@@ -6,6 +6,7 @@
 import shutil
 import filecmp
 import subprocess
+import signal
 
 from dvc.main import main
 from dvc.output import OutputBase
@@ -323,16 +324,48 @@ def test_run_args_with_spaces(self):
         self.assertEqual(ret, 0)
         self.assertEqual(stage.cmd, 'echo "foo bar"')
 
-    @mock.patch.object(subprocess, "Popen", side_effect=KeyboardInterrupt)
-    def test_keyboard_interrupt_before_communicate(self, _):
-        ret = main(["run", "mycmd"])
-        self.assertEqual(ret, 252)
-
     @mock.patch.object(subprocess.Popen, "wait", new=KeyboardInterrupt)
-    def test_keyboard_interrupt_during_communicate(self):
-        ret = main(["run", "python", "code.py"])
+    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)
+
 
 class TestRunRemoveOuts(TestDvc):
     def test(self):