Merge pull request #27285 from anntzer/ai

ksunden · web-flow · commit aeefb3b4cd9f · 2023-11-15T14:08:47.000-06:00
Factor out common parts of qt and macos interrupt handling.
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)
