diff --git a/dvc/stage.py b/dvc/stage.py index 93831657e1..7610653f35 100644 --- 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 @@ -767,16 +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, - ) - 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): diff --git a/tests/func/test_run.py b/tests/func/test_run.py index 79ae192db0..3d707cdda5 100644 --- a/tests/func/test_run.py +++ b/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,9 +324,46 @@ 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(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)