Always raise, don't swallow exception, if exception is `KeyboardInter…

…rupt` in `ProtocolUnit.execute` (#215)

* Always raise, don't swallow exception, if exception is `KeyboardInterrupt` in `ProtocolUnit.execute`

* Added explicit test for KeyboardInterrupt raised in ProtocolUnit._execute

---------

Co-authored-by: Mike Henry <[email protected]>

gufe/protocols/protocolunit.py

@@ -323,6 +323,9 @@ def execute(self, *,
                 start_time=start, end_time=datetime.datetime.now(),
             )
 
+        except KeyboardInterrupt:
+            # if we "fail" due to a KeyboardInterrupt, we always want to raise
+            raise
         except Exception as e:
             if raise_error:
                 raise

gufe/tests/test_protocolunit.py

@@ -2,7 +2,7 @@
 import pytest
 from pathlib import Path
 
-from gufe.protocols.protocolunit import ProtocolUnit, Context, ProtocolUnitFailure
+from gufe.protocols.protocolunit import ProtocolUnit, Context, ProtocolUnitFailure, ProtocolUnitResult
 from gufe.tests.test_tokenization import GufeTokenizableTestsMixin
 
 
@@ -16,6 +16,16 @@ def _execute(ctx: Context, an_input=2, **inputs):
         return {"foo": "bar"}
 
 
+class DummyKeyboardInterruptUnit(ProtocolUnit):
+    @staticmethod
+    def _execute(ctx: Context, an_input=2, **inputs):
+
+        if an_input != 2:
+            raise KeyboardInterrupt
+
+        return {"foo": "bar"}
+
+
 @pytest.fixture
 def dummy_unit():
     return DummyUnit(name="qux")
@@ -66,6 +76,26 @@ def test_execute(self, tmpdir):
             with pytest.raises(ValueError, match="should always be 2"):
                 unit.execute(context=ctx, raise_error=True, an_input=3)
 
+    def test_execute_KeyboardInterrupt(self, tmpdir):
+        with tmpdir.as_cwd():
+
+            unit = DummyKeyboardInterruptUnit()
+
+            shared = Path('shared') / str(unit.key)
+            shared.mkdir(parents=True)
+
+            scratch = Path('scratch') / str(unit.key)
+            scratch.mkdir(parents=True)
+
+            ctx = Context(shared=shared, scratch=scratch)
+
+            with pytest.raises(KeyboardInterrupt):
+                unit.execute(context=ctx, an_input=3)
+
+            u: ProtocolUnitResult = unit.execute(context=ctx, an_input=2)
+
+            assert u.outputs == {'foo': 'bar'}
+
     def test_normalize(self, dummy_unit):
         thingy = dummy_unit.key