diff --git a/dvc/stage.py b/dvc/stage.py index 3f7290f330..0174a2f00f 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 @@ -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 index 79ae192db0..6365f04bdf 100644 --- 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):