Create context manager that allows a loop to be interrupted by pressi…

…ng Ctrl+c

docs/source/_templates/api_reference_class_template.rst

@@ -5,6 +5,7 @@
    ConstantUnitMathConventions
    dev
    exps
+   InterruptibleLoop
    ndarray
    np
    num

docs/source/api_reference/dev.rst

@@ -16,4 +16,5 @@ profiling code.
     :toctree: ./api
     :template: ../_templates/api_reference_class_template.rst
 
+    InterruptibleLoop
     TimeIt

docs/source/spelling_wordlist.txt

@@ -5,6 +5,7 @@ gigajoule
 gigajoules
 gigapascal
 gigapascals
+interruptible
 iterable
 iteratively
 Iteratively

pyxx/dev/__init__.py

@@ -4,4 +4,4 @@
 or scripts that can be useful when debugging Python code.
 """
 
-from .classes import TimeIt
+from .classes import InterruptibleLoop, TimeIt

pyxx/dev/classes/__init__.py

@@ -2,4 +2,5 @@
 Python code.
 """
 
+from .loops import InterruptibleLoop
 from .timer import TimeIt

pyxx/dev/classes/loops.py

@@ -0,0 +1,71 @@
+"""Classes for controlling loop execution.
+"""
+
+import signal
+
+
+class InterruptibleLoop:
+    """Context manager to create an interruptible loop
+
+    This context manager allows users to interrupt and terminate a loop early
+    by pressing ``Ctrl+c``.  This can be useful to stop a long process early
+    while monitoring results (such as stopping training of a machine learning
+    model).  When interrupted, the context manager will change its
+    :py:attr:`interrupted` attribute to ``True``.  Additionally, if the
+    ``throw_exception`` option is enabled, an ``InterruptedError`` will be
+    raised (which can be caught and handled as desired).
+
+    Examples
+    --------
+    Basic usage:
+
+    >>> with pyxx.dev.InterruptibleLoop() as loop:
+    ...     for i in range(1000):
+    ...         # Do something
+    ...
+    ...         if loop.interrupted:
+    ...             print('Loop was interrupted')
+    ...             break
+
+    Example of using the ``throw_exception`` option:
+
+    >>> try:
+    ...     with pyxx.dev.InterruptibleLoop(throw_exception=True) as loop:
+    ...         for i in range(1000):
+    ...             pass  # Do something
+    ...
+    ... except InterruptedError:
+    ...     print('Loop was interrupted')
+    ...
+    ... else:
+    ...     print('Loop was NOT interrupted')
+    Loop was NOT interrupted
+    """
+
+    def __init__(self, throw_exception: bool = False) -> None:
+        # Store user inputs
+        self.__throw_exception = throw_exception
+
+        # Initialize context manager internal state
+        self.__interrupted = False
+        self.__original_sigint_handler = signal.getsignal(signal.SIGINT)
+
+    @property
+    def interrupted(self) -> bool:
+        """Whether the loop has been interrupted"""
+        return self.__interrupted
+
+    def _interrupt_handler(self, *args, **kwargs) -> None:  # pylint: disable=W0613
+        self.__interrupted = True
+
+        if self.__throw_exception:
+            raise InterruptedError
+
+    def __enter__(self) -> 'InterruptibleLoop':
+        signal.signal(signal.SIGINT, self._interrupt_handler)
+        self.__interrupted = False
+
+        return self
+
+    def __exit__(self, *args, **kwargs) -> None:
+        signal.signal(signal.SIGINT, self.__original_sigint_handler)

tests/dev/classes/__init__.py

@@ -1 +1,2 @@
+from .test_loops import *
 from .test_timer import *

tests/dev/classes/test_loops.py

@@ -0,0 +1,34 @@
+import signal
+import unittest
+
+from pyxx.dev import InterruptibleLoop
+
+
+class Test_InterruptibleLoop(unittest.TestCase):
+    def test_interrupted_attribute(self):
+        # Verifies that the `InterruptibleLoop.interrupted` attribute
+        # changes when Ctrl+c is pressed
+        with InterruptibleLoop() as loop:
+            self.assertFalse(loop.interrupted)
+
+            signal.raise_signal(signal.SIGINT)
+            self.assertTrue(loop.interrupted)
+
+    def test_restore_handler(self):
+        # Verifies that the `InterruptibleLoop` class restores the
+        # original SIGINT handler when the context manager exits
+        sigint_handler = signal.getsignal(signal.SIGINT)
+
+        self.assertIs(signal.getsignal(signal.SIGINT), sigint_handler)
+
+        with InterruptibleLoop():
+            self.assertIsNot(signal.getsignal(signal.SIGINT), sigint_handler)
+
+        self.assertIs(signal.getsignal(signal.SIGINT), sigint_handler)
+
+    def test_throw_exception(self):
+        # Verifies that an exception is thrown when pressing Ctrl+c
+        # if the `throw_exception` option is set to `True`
+        with InterruptibleLoop(throw_exception=True):
+            with self.assertRaises(InterruptedError):
+                signal.raise_signal(signal.SIGINT)