replace exception catch with signal

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
@@ -775,10 +776,9 @@ class Stage(object):
             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)

tests/func/test_run.py

@@ -324,10 +324,15 @@ class TestCmdRun(TestDvc):
         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):