Merge pull request microsoft#6192 from trevormorgan/trmorgan/keyboard…

…-interrupt

Trmorgan/keyboard interrupt

docs/changes/0.46.0.rst

@@ -27,6 +27,10 @@ Breaking Changes:
   If the attribute is not a ParameterBase this will instead warn. It is the intention that this becomes an error in the future.
   (:pr:`6174`) (:pr:`6211`)
 
+- Updated dond functions to to re-raise KeyboardInterrupt for better interrupt handling making it easier to stop long-running measurement
+  loops and reducing the need for kernel restarts. This meas that if you interrupt a `dond`` function with a keyboard interrupt not only
+  the measurement but any pending code to execute will be interrupted. In the process logging for interrupted measurements has been improved. (:pr:`6192`)
+
 
 Improved:
 ---------

src/qcodes/dataset/dond/do_nd_utils.py

@@ -120,14 +120,26 @@ def _register_actions(
 
 
 @contextmanager
-def catch_interrupts() -> Iterator[Callable[[], MeasInterruptT]]:
-    interrupt_exception = None
+def catch_interrupts() -> Iterator[Callable[[], MeasInterruptT | None]]:
+    interrupt_exception: MeasInterruptT | None = None
+    interrupt_raised = False
 
-    def get_interrupt_exception() -> MeasInterruptT:
+    def get_interrupt_exception() -> MeasInterruptT | None:
         nonlocal interrupt_exception
         return interrupt_exception
 
     try:
         yield get_interrupt_exception
-    except (KeyboardInterrupt, BreakConditionInterrupt) as e:
+    except KeyboardInterrupt as e:
         interrupt_exception = e
+        interrupt_raised = True
+        raise  # Re-raise KeyboardInterrupt
+    except BreakConditionInterrupt as e:
+        interrupt_exception = e
+        interrupt_raised = True
+        # Don't re-raise BreakConditionInterrupt
+    finally:
+        if interrupt_raised:
+            log.warning(
+                f"Measurement has been interrupted, data may be incomplete: {interrupt_exception}"
+            )

tests/dataset/dond/test_dond_keyboard_interrupts.py

@@ -0,0 +1,100 @@
+
+import pytest
+
+from qcodes.dataset.dond.do_nd_utils import BreakConditionInterrupt, catch_interrupts
+
+
+def test_catch_interrupts():
+    # Test normal execution (no interrupt)
+    with catch_interrupts() as get_interrupt:
+        assert get_interrupt() is None
+
+    # Test KeyboardInterrupt
+    with pytest.raises(KeyboardInterrupt):
+        with catch_interrupts() as get_interrupt:
+            raise KeyboardInterrupt()
+
+    # Test BreakConditionInterrupt
+    with catch_interrupts() as get_interrupt:
+        raise BreakConditionInterrupt()
+    assert isinstance(get_interrupt(), BreakConditionInterrupt)
+
+    # Test that cleanup code runs for KeyboardInterrupt
+    cleanup_ran = False
+    with pytest.raises(KeyboardInterrupt):
+        with catch_interrupts():
+            try:
+                raise KeyboardInterrupt()
+            finally:
+                cleanup_ran = True
+    assert cleanup_ran
+
+    # Test that cleanup code runs for BreakConditionInterrupt
+    cleanup_ran = False
+    with catch_interrupts():
+        try:
+            raise BreakConditionInterrupt()
+        finally:
+            cleanup_ran = True
+    assert cleanup_ran
+
+    # Test that BreakConditionInterrupt is caught and doesn't raise
+    with catch_interrupts() as get_interrupt:
+        raise BreakConditionInterrupt()
+    assert isinstance(get_interrupt(), BreakConditionInterrupt)
+
+
+def test_catch_interrupts_in_loops():
+    # Test interruption in a simple loop
+    loop_count = 0
+    with pytest.raises(KeyboardInterrupt):
+        for i in range(5):
+            with catch_interrupts():
+                loop_count += 1
+                if i == 2:
+                    raise KeyboardInterrupt()
+    assert loop_count == 3  # Loop should stop at the third iteration
+
+    # Test interruption in nested loops
+    outer_count = 0
+    inner_count = 0
+    with pytest.raises(KeyboardInterrupt):
+        for i in range(3):
+            with catch_interrupts():
+                outer_count += 1
+                for j in range(3):
+                    with catch_interrupts():
+                        inner_count += 1
+                        if i == 1 and j == 1:
+                            raise KeyboardInterrupt()
+    assert outer_count == 2
+    assert inner_count == 5
+
+
+def test_catch_interrupts_simulated_sweeps():
+    def simulated_sweep(interrupt_at=None):
+        for i in range(5):
+            with catch_interrupts():
+                if i == interrupt_at:
+                    raise KeyboardInterrupt()
+                yield i
+
+    # Test interruption in a single sweep
+    results = []
+    with pytest.raises(KeyboardInterrupt):
+        for value in simulated_sweep(interrupt_at=3):
+            results.append(value)
+    assert results == [0, 1, 2]
+
+    # Test interruption in nested sweeps
+    outer_results = []
+    inner_results = []
+    with pytest.raises(KeyboardInterrupt):
+        for outer_value in simulated_sweep(interrupt_at=None):
+            outer_results.append(outer_value)
+            for inner_value in simulated_sweep(
+                interrupt_at=2 if outer_value == 1 else None
+            ):
+                inner_results.append(inner_value)
+    assert outer_results == [0, 1]
+    assert inner_results == [0, 1, 2, 3, 4, 0, 1]