replace exception catch with signal

efiop · Jun 20, 2019 · 5457c16 · 5457c16
1 parent ccb8e33
commit 5457c16
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/dvc/stage.py b/dvc/stage.py
@@ -7,6 +7,7 @@
 import os
 import subprocess
 import logging
+import signal
 
 from dvc.utils import relpath
 from dvc.utils.compat import pathlib
@@ -775,10 +776,9 @@ def _run(self):
             executable=executable,
         )
 
-        try:
-            p.communicate()
-        except KeyboardInterrupt:
-            p.communicate()
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
+        p.communicate()
+        signal.signal(signal.SIGINT, signal.default_int_handler)
 
         if p.returncode != 0:
             raise StageCmdFailedError(self)

diff --git a/tests/func/test_run.py b/tests/func/test_run.py
@@ -324,10 +324,15 @@ def test_run_args_with_spaces(self):
         self.assertEqual(stage.cmd, 'echo "foo bar"')
 
     @mock.patch.object(subprocess, "Popen", side_effect=KeyboardInterrupt)
-    def test_keyboard_interrupt(self, _):
+    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"])
+        self.assertEqual(ret, 1)
+
 
 class TestRunRemoveOuts(TestDvc):
     def test(self):