From ccb8e338e1772778ea871e02e4b0b0f90ee2a234 Mon Sep 17 00:00:00 2001
From: Kirill Vasin <k.vasin@semrush.com>
Date: Tue, 18 Jun 2019 19:11:49 +0300
Subject: [PATCH 1/3] added custom KeyboardInterrupt support

---
 dvc/stage.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/dvc/stage.py b/dvc/stage.py
index 93831657e1..3f7290f330 100644
--- a/dvc/stage.py
+++ b/dvc/stage.py
@@ -774,7 +774,11 @@ def _run(self):
             env=fix_env(os.environ),
             executable=executable,
         )
-        p.communicate()
+
+        try:
+            p.communicate()
+        except KeyboardInterrupt:
+            p.communicate()
 
         if p.returncode != 0:
             raise StageCmdFailedError(self)

From 5457c16e766000822ae31760bb29e97457cf2d28 Mon Sep 17 00:00:00 2001
From: Kirill Vasin <k.vasin@semrush.com>
Date: Thu, 20 Jun 2019 13:39:22 +0300
Subject: [PATCH 2/3] replace exception catch with signal

---
 dvc/stage.py           | 8 ++++----
 tests/func/test_run.py | 7 ++++++-
 2 files changed, 10 insertions(+), 5 deletions(-)

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):

From 1734e25827180e4e82255eea0a29a2fabfc87be4 Mon Sep 17 00:00:00 2001
From: Kirill Vasin <k.vasin@semrush.com>
Date: Thu, 20 Jun 2019 19:21:53 +0300
Subject: [PATCH 3/3] added old signal persistance and fixed tests accordingly

---
 dvc/stage.py           | 25 ++++++++++++----------
 tests/func/test_run.py | 47 +++++++++++++++++++++++++++++++++++-------
 2 files changed, 54 insertions(+), 18 deletions(-)

diff --git a/dvc/stage.py b/dvc/stage.py
index 0174a2f00f..7610653f35 100644
--- a/dvc/stage.py
+++ b/dvc/stage.py
@@ -768,19 +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,
-        )
+        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 6365f04bdf..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,16 +324,48 @@ 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_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):