diff --git a/dvc/stage.py b/dvc/stage.py index 0174a2f00..7610653f3 100644 --- a/dvc/stage.py +++ b/dvc/stage.py @@ -768,19 +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, - ) + 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): diff --git a/tests/func/test_run.py b/tests/func/test_run.py index 6365f04bd..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,16 +324,48 @@ 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_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):