Factor out common parts of qt and macos interrupt handling.

anntzer · anntzer · commit 3bd9abe6b366 · 2023-11-08T12:54:27.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/backend_bases.py b/lib/matplotlib/backend_bases.py
@@ -35,6 +35,8 @@
 import itertools
 import logging
 import os
+import signal
+import socket
 import sys
 import time
 import weakref
@@ -1651,6 +1653,64 @@ def _is_non_interactive_terminal_ipython(ip):
             and getattr(ip.parent, 'interact', None) is False)
 
 
+@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 the Python interpreter
+    from running and processing signals (i.e., to raise a KeyboardInterrupt).
+    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)
+
+
 class FigureCanvasBase:
     """
     The canvas the figure renders into.
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
@@ -10,14 +7,20 @@
 from .backend_agg import FigureCanvasAgg
 from matplotlib.backend_bases import (
     _Backend, FigureCanvasBase, FigureManagerBase, NavigationToolbar2,
-    ResizeEvent, TimerBase)
+    ResizeEvent, TimerBase, _allow_interrupt)
 
 
 class TimerMac(_macosx.Timer, TimerBase):
     """Subclass of `.TimerBase` using CFRunLoop timer events."""
     # 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 _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
@@ -9,13 +9,12 @@
 from matplotlib.backend_bases import (
     _Backend, FigureCanvasBase, FigureManagerBase, NavigationToolbar2,
     TimerBase, cursors, ToolContainerBase, MouseButton,
-    CloseEvent, KeyEvent, LocationEvent, MouseEvent, ResizeEvent)
+    CloseEvent, KeyEvent, LocationEvent, MouseEvent, ResizeEvent,
+    _allow_interrupt)
 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 +147,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 _allow_interrupt(prepare_notifier, handle_sigint)
+
+
 class TimerQT(TimerBase):
     """Subclass of `.TimerBase` using QTimer events."""
 
@@ -417,7 +448,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 +629,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)
