Factor out common parts of qt and macos interrupt handling.

anntzer · anntzer · commit 29de7a168d94 · 2023-11-07T23:47:50.000+01:00
Note that we don't actually need to disable the QSocketNotifier at the
end, just letting it go out of scope should be sufficient as its
destructor also does that (see qsocketnotifier.cpp).
diff --git a/lib/matplotlib/backends/backend_macosx.py b/lib/matplotlib/backends/backend_macosx.py
@@ -1,7 +1,4 @@
-import contextlib
 import os
-import signal
-import socket
 
 import matplotlib as mpl
 from matplotlib import _api, cbook
@@ -18,6 +15,12 @@ class TimerMac(_macosx.Timer, TimerBase):
     # completely implemented at the C-level (in _macosx.Timer)
 
 
+def _allow_interrupt_macos():
+    """A context manager that allows terminating a plot by sending a SIGINT."""
+    return cbook._allow_interrupt(
+        lambda rsock: _macosx.wake_on_fd_write(rsock.fileno()), _macosx.stop)
+
+
 class FigureCanvasMac(FigureCanvasAgg, _macosx.FigureCanvas, FigureCanvasBase):
     # docstring inherited
 
@@ -109,10 +112,9 @@ def resize(self, width, height):
 
     def start_event_loop(self, timeout=0):
         # docstring inherited
-        with _maybe_allow_interrupt():
-            # Call the objc implementation of the event loop after
-            # setting up the interrupt handling
-            self._start_event_loop(timeout=timeout)
+        # Set up a SIGINT handler to allow terminating a plot via CTRL-C.
+        with _allow_interrupt_macos():
+            self._start_event_loop(timeout=timeout)  # Forward to ObjC implementation.
 
 
 class NavigationToolbar2Mac(_macosx.NavigationToolbar2, NavigationToolbar2):
@@ -177,9 +179,7 @@ def destroy(self):
     @classmethod
     def start_main_loop(cls):
         # Set up a SIGINT handler to allow terminating a plot via CTRL-C.
-        # The logic is largely copied from qt_compat._maybe_allow_interrupt; see its
-        # docstring for details.  Parts are implemented by wake_on_fd_write in ObjC.
-        with _maybe_allow_interrupt():
+        with _allow_interrupt_macos():
             _macosx.show()
 
     def show(self):
@@ -190,45 +190,6 @@ def show(self):
             self._raise()
 
 
-@contextlib.contextmanager
-def _maybe_allow_interrupt():
-    """
-    This manager allows to terminate a plot by sending a SIGINT. It is
-    necessary because the running backend prevents Python interpreter to
-    run and process signals (i.e., to raise KeyboardInterrupt exception). To
-    solve this one needs to somehow wake up the interpreter and make it close
-    the plot window. The implementation is taken from qt_compat, see that
-    docstring for a more detailed description.
-    """
-    old_sigint_handler = signal.getsignal(signal.SIGINT)
-    if old_sigint_handler in (None, signal.SIG_IGN, signal.SIG_DFL):
-        yield
-        return
-
-    handler_args = None
-    wsock, rsock = socket.socketpair()
-    wsock.setblocking(False)
-    rsock.setblocking(False)
-    old_wakeup_fd = signal.set_wakeup_fd(wsock.fileno())
-    _macosx.wake_on_fd_write(rsock.fileno())
-
-    def handle(*args):
-        nonlocal handler_args
-        handler_args = args
-        _macosx.stop()
-
-    signal.signal(signal.SIGINT, handle)
-    try:
-        yield
-    finally:
-        wsock.close()
-        rsock.close()
-        signal.set_wakeup_fd(old_wakeup_fd)
-        signal.signal(signal.SIGINT, old_sigint_handler)
-        if handler_args is not None:
-            old_sigint_handler(*handler_args)
-
-
 @_Backend.export
 class _BackendMac(_Backend):
     FigureCanvas = FigureCanvasMac
diff --git a/lib/matplotlib/backends/backend_qt.py b/lib/matplotlib/backends/backend_qt.py
@@ -13,9 +13,7 @@
 import matplotlib.backends.qt_editor.figureoptions as figureoptions
 from . import qt_compat
 from .qt_compat import (
-    QtCore, QtGui, QtWidgets, __version__, QT_API,
-    _to_int, _isdeleted, _maybe_allow_interrupt
-)
+    QtCore, QtGui, QtWidgets, __version__, QT_API, _to_int, _isdeleted)
 
 
 # SPECIAL_KEYS are Qt::Key that do *not* return their Unicode name
@@ -148,6 +146,38 @@ def _create_qApp():
     return app
 
 
+def _allow_interrupt_qt(qapp_or_eventloop):
+    """A context manager that allows terminating a plot by sending a SIGINT."""
+
+    # Use QSocketNotifier to read the socketpair while the Qt event loop runs.
+
+    def prepare_notifier(rsock):
+        sn = QtCore.QSocketNotifier(rsock.fileno(), QtCore.QSocketNotifier.Type.Read)
+
+        @sn.activated.connect
+        def _may_clear_sock():
+            # Running a Python function on socket activation gives the interpreter a
+            # chance to handle the signal in Python land.  We also need to drain the
+            # socket with recv() to re-arm it, because it will be written to as part of
+            # the wakeup.  (We need this in case set_wakeup_fd catches a signal other
+            # than SIGINT and we shall continue waiting.)
+            try:
+                rsock.recv(1)
+            except BlockingIOError:
+                # This may occasionally fire too soon or more than once on Windows, so
+                # be forgiving about reading an empty socket.
+                pass
+
+        return sn  # Actually keep the notifier alive.
+
+    def handle_sigint():
+        if hasattr(qapp_or_eventloop, 'closeAllWindows'):
+            qapp_or_eventloop.closeAllWindows()
+        qapp_or_eventloop.quit()
+
+    return mpl.cbook._allow_interrupt(prepare_notifier, handle_sigint)
+
+
 class TimerQT(TimerBase):
     """Subclass of `.TimerBase` using QTimer events."""
 
@@ -417,7 +447,7 @@ def start_event_loop(self, timeout=0):
         if timeout > 0:
             _ = QtCore.QTimer.singleShot(int(timeout * 1000), event_loop.quit)
 
-        with _maybe_allow_interrupt(event_loop):
+        with _allow_interrupt_qt(event_loop):
             qt_compat._exec(event_loop)
 
     def stop_event_loop(self, event=None):
@@ -598,7 +628,7 @@ def resize(self, width, height):
     def start_main_loop(cls):
         qapp = QtWidgets.QApplication.instance()
         if qapp:
-            with _maybe_allow_interrupt(qapp):
+            with _allow_interrupt_qt(qapp):
                 qt_compat._exec(qapp)
 
     def show(self):
diff --git a/lib/matplotlib/backends/qt_compat.py b/lib/matplotlib/backends/qt_compat.py
@@ -13,9 +13,6 @@
 import os
 import platform
 import sys
-import signal
-import socket
-import contextlib
 
 from packaging.version import parse as parse_version
 
@@ -160,73 +157,3 @@ def _isdeleted(obj):
 def _exec(obj):
     # exec on PyQt6, exec_ elsewhere.
     obj.exec() if hasattr(obj, "exec") else obj.exec_()
-
-
-@contextlib.contextmanager
-def _maybe_allow_interrupt(qapp_or_eventloop):
-    """
-    This manager allows to terminate a plot by sending a SIGINT. It is
-    necessary because the running Qt backend prevents Python interpreter to
-    run and process signals (i.e., to raise KeyboardInterrupt exception). To
-    solve this one needs to somehow wake up the interpreter and make it close
-    the plot window. We do this by using the signal.set_wakeup_fd() function
-    which organizes a write of the signal number into a socketpair connected
-    to the QSocketNotifier (since it is part of the Qt backend, it can react
-    to that write event). Afterwards, the Qt handler empties the socketpair
-    by a recv() command to re-arm it (we need this if a signal different from
-    SIGINT was caught by set_wakeup_fd() and we shall continue waiting). If
-    the SIGINT was caught indeed, after exiting the on_signal() function the
-    interpreter reacts to the SIGINT according to the handle() function which
-    had been set up by a signal.signal() call: it causes the qt_object to
-    exit by calling its quit() method. Finally, we call the old SIGINT
-    handler with the same arguments that were given to our custom handle()
-    handler.
-
-    We do this only if the old handler for SIGINT was not None, which means
-    that a non-python handler was installed, i.e. in Julia, and not SIG_IGN
-    which means we should ignore the interrupts.
-    """
-
-    old_sigint_handler = signal.getsignal(signal.SIGINT)
-    if old_sigint_handler in (None, signal.SIG_IGN, signal.SIG_DFL):
-        yield
-        return
-
-    handler_args = None
-    wsock, rsock = socket.socketpair()
-    wsock.setblocking(False)
-    rsock.setblocking(False)
-    old_wakeup_fd = signal.set_wakeup_fd(wsock.fileno())
-    sn = QtCore.QSocketNotifier(rsock.fileno(), QtCore.QSocketNotifier.Type.Read)
-
-    # We do not actually care about this value other than running some Python code to
-    # ensure that the interpreter has a chance to handle the signal in Python land.  We
-    # also need to drain the socket because it will be written to as part of the wakeup!
-    # There are some cases where this may fire too soon / more than once on Windows so
-    # we should be forgiving about reading an empty socket.
-    # Clear the socket to re-arm the notifier.
-    @sn.activated.connect
-    def _may_clear_sock(*args):
-        try:
-            rsock.recv(1)
-        except BlockingIOError:
-            pass
-
-    def handle(*args):
-        nonlocal handler_args
-        handler_args = args
-        if hasattr(qapp_or_eventloop, 'closeAllWindows'):
-            qapp_or_eventloop.closeAllWindows()
-        qapp_or_eventloop.quit()
-
-    signal.signal(signal.SIGINT, handle)
-    try:
-        yield
-    finally:
-        wsock.close()
-        rsock.close()
-        sn.setEnabled(False)
-        signal.set_wakeup_fd(old_wakeup_fd)
-        signal.signal(signal.SIGINT, old_sigint_handler)
-        if handler_args is not None:
-            old_sigint_handler(*handler_args)
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -14,6 +14,8 @@
 import os
 from pathlib import Path
 import shlex
+import signal
+import socket
 import subprocess
 import sys
 import time
@@ -82,6 +84,64 @@ def _get_running_interactive_framework():
     return None
 
 
+@contextlib.contextmanager
+def _allow_interrupt(prepare_notifier, handle_sigint):
+    """
+    A context manager that allows terminating a plot by sending a SIGINT.  It
+    is necessary because the running backend prevents Python interpreter to
+    run and process signals (i.e., to raise KeyboardInterrupt exception).  To
+    solve this one needs to somehow wake up the interpreter and make it close
+    the plot window.  We do this by using the signal.set_wakeup_fd() function
+    which organizes a write of the signal number into a socketpair.  A
+    backend-specific function, *prepare_notifier*, arranges to listen to the
+    pair's read socket while the event loop is running.  (If it returns a
+    notifier object, that object is kept alive while the context manager runs.)
+
+    If SIGINT was indeed caught, after exiting the on_signal() function the
+    interpreter reacts to the signal according to the handler function which
+    had been set up by a signal.signal() call; here, we arrange to call the
+    backend-specific *handle_sigint* function.  Finally, we call the old SIGINT
+    handler with the same arguments that were given to our custom handler.
+
+    We do this only if the old handler for SIGINT was not None, which means
+    that a non-python handler was installed, i.e. in Julia, and not SIG_IGN
+    which means we should ignore the interrupts.
+
+    Parameters
+    ----------
+    prepare_notifier : Callable[[socket.socket], object]
+    handle_sigint : Callable[[], object]
+    """
+
+    old_sigint_handler = signal.getsignal(signal.SIGINT)
+    if old_sigint_handler in (None, signal.SIG_IGN, signal.SIG_DFL):
+        yield
+        return
+
+    handler_args = None
+    wsock, rsock = socket.socketpair()
+    wsock.setblocking(False)
+    rsock.setblocking(False)
+    old_wakeup_fd = signal.set_wakeup_fd(wsock.fileno())
+    notifier = prepare_notifier(rsock)
+
+    def save_args_and_handle_sigint(*args):
+        nonlocal handler_args
+        handler_args = args
+        handle_sigint()
+
+    signal.signal(signal.SIGINT, save_args_and_handle_sigint)
+    try:
+        yield
+    finally:
+        wsock.close()
+        rsock.close()
+        signal.set_wakeup_fd(old_wakeup_fd)
+        signal.signal(signal.SIGINT, old_sigint_handler)
+        if handler_args is not None:
+            old_sigint_handler(*handler_args)
+
+
 def _exception_printer(exc):
     if _get_running_interactive_framework() in ["headless", None]:
         raise exc
