Merge pull request #13306 from vdrhtc/qt5-sigint

QuLogic · web-flow · commit 9e8c8faa1b2e · 2021-08-23T18:13:33.000-04:00
Qt5: SIGINT kills just the mpl window and not the process itself
diff --git a/lib/matplotlib/backends/backend_qt.py b/lib/matplotlib/backends/backend_qt.py
@@ -17,8 +17,10 @@
     QtCore, QtGui, QtWidgets, __version__, QT_API,
     _enum, _to_int,
     _devicePixelRatioF, _isdeleted, _setDevicePixelRatio,
+    _maybe_allow_interrupt
 )
 
+
 backend_version = __version__
 
 # SPECIAL_KEYS are Qt::Key that do *not* return their unicode name
@@ -399,7 +401,9 @@ def start_event_loop(self, timeout=0):
         if timeout > 0:
             timer = QtCore.QTimer.singleShot(int(timeout * 1000),
                                              event_loop.quit)
-        qt_compat._exec(event_loop)
+
+        with _maybe_allow_interrupt(event_loop):
+            qt_compat._exec(event_loop)
 
     def stop_event_loop(self, event=None):
         # docstring inherited
@@ -1001,14 +1005,5 @@ class _BackendQT(_Backend):
 
     @staticmethod
     def mainloop():
-        old_signal = signal.getsignal(signal.SIGINT)
-        # allow SIGINT exceptions to close the plot window.
-        is_python_signal_handler = old_signal is not None
-        if is_python_signal_handler:
-            signal.signal(signal.SIGINT, signal.SIG_DFL)
-        try:
+        with _maybe_allow_interrupt(qApp):
             qt_compat._exec(qApp)
-        finally:
-            # reset the SIGINT exception handler
-            if is_python_signal_handler:
-                signal.signal(signal.SIGINT, old_signal)
diff --git a/lib/matplotlib/backends/qt_compat.py b/lib/matplotlib/backends/qt_compat.py
@@ -16,6 +16,9 @@
 import os
 import platform
 import sys
+import signal
+import socket
+import contextlib
 
 from packaging.version import parse as parse_version
 
@@ -28,7 +31,7 @@
 QT_API_PYSIDE2 = "PySide2"
 QT_API_PYQTv2 = "PyQt4v2"
 QT_API_PYSIDE = "PySide"
-QT_API_PYQT = "PyQt4"   # Use the old sip v1 API (Py3 defaults to v2).
+QT_API_PYQT = "PyQt4"  # Use the old sip v1 API (Py3 defaults to v2).
 QT_API_ENV = os.environ.get("QT_API")
 if QT_API_ENV is not None:
     QT_API_ENV = QT_API_ENV.lower()
@@ -74,32 +77,33 @@
 
 
 def _setup_pyqt5plus():
-    global QtCore, QtGui, QtWidgets, __version__, is_pyqt5, \
+    global QtCore, QtGui, QtWidgets, QtNetwork, __version__, is_pyqt5, \
         _isdeleted, _getSaveFileName
 
     if QT_API == QT_API_PYQT6:
-        from PyQt6 import QtCore, QtGui, QtWidgets, sip
+        from PyQt6 import QtCore, QtGui, QtWidgets, QtNetwork, sip
         __version__ = QtCore.PYQT_VERSION_STR
         QtCore.Signal = QtCore.pyqtSignal
         QtCore.Slot = QtCore.pyqtSlot
         QtCore.Property = QtCore.pyqtProperty
         _isdeleted = sip.isdeleted
     elif QT_API == QT_API_PYSIDE6:
-        from PySide6 import QtCore, QtGui, QtWidgets, __version__
+        from PySide6 import QtCore, QtGui, QtWidgets, QtNetwork, __version__
         import shiboken6
         def _isdeleted(obj): return not shiboken6.isValid(obj)
     elif QT_API == QT_API_PYQT5:
-        from PyQt5 import QtCore, QtGui, QtWidgets
+        from PyQt5 import QtCore, QtGui, QtWidgets, QtNetwork
         import sip
         __version__ = QtCore.PYQT_VERSION_STR
         QtCore.Signal = QtCore.pyqtSignal
         QtCore.Slot = QtCore.pyqtSlot
         QtCore.Property = QtCore.pyqtProperty
         _isdeleted = sip.isdeleted
     elif QT_API == QT_API_PYSIDE2:
-        from PySide2 import QtCore, QtGui, QtWidgets, __version__
+        from PySide2 import QtCore, QtGui, QtWidgets, QtNetwork, __version__
         import shiboken2
-        def _isdeleted(obj): return not shiboken2.isValid(obj)
+        def _isdeleted(obj):
+            return not shiboken2.isValid(obj)
     else:
         raise AssertionError(f"Unexpected QT_API: {QT_API}")
     _getSaveFileName = QtWidgets.QFileDialog.getSaveFileName
@@ -134,7 +138,6 @@ def _isdeleted(obj): return not shiboken2.isValid(obj)
         "QT_MAC_WANTS_LAYER" not in os.environ):
     os.environ["QT_MAC_WANTS_LAYER"] = "1"
 
-
 # These globals are only defined for backcompatibility purposes.
 ETS = dict(pyqt5=(QT_API_PYQT5, 5), pyside2=(QT_API_PYSIDE2, 5))
 
@@ -191,3 +194,62 @@ def _setDevicePixelRatio(obj, val):
     if hasattr(obj, 'setDevicePixelRatio'):
         # Not available on Qt4 or some older Qt5.
         obj.setDevicePixelRatio(val)
+
+
+@contextlib.contextmanager
+def _maybe_allow_interrupt(qapp):
+    """
+    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)
+    handler_args = None
+    skip = False
+    if old_sigint_handler in (None, signal.SIG_IGN, signal.SIG_DFL):
+        skip = True
+    else:
+        wsock, rsock = socket.socketpair()
+        wsock.setblocking(False)
+        old_wakeup_fd = signal.set_wakeup_fd(wsock.fileno())
+        sn = QtCore.QSocketNotifier(
+            rsock.fileno(), _enum('QtCore.QSocketNotifier.Type').Read
+        )
+
+        # Clear the socket to re-arm the notifier.
+        sn.activated.connect(lambda *args: rsock.recv(1))
+
+        def handle(*args):
+            nonlocal handler_args
+            handler_args = args
+            qapp.quit()
+
+        signal.signal(signal.SIGINT, handle)
+    try:
+        yield
+    finally:
+        if not skip:
+            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/tests/test_backend_qt.py b/lib/matplotlib/tests/test_backend_qt.py
@@ -33,8 +33,22 @@ def qt_core(request):
     return QtCore
 
 
+@pytest.fixture
+def platform_simulate_ctrl_c(request):
+    import signal
+    from functools import partial
+
+    if hasattr(signal, "CTRL_C_EVENT"):
+        from win32api import GenerateConsoleCtrlEvent
+        return partial(GenerateConsoleCtrlEvent, 0, 0)
+    else:
+        # we're not on windows
+        return partial(os.kill, os.getpid(), signal.SIGINT)
+
+
 @pytest.mark.backend('QtAgg', skip_on_importerror=True)
 def test_fig_close():
+
     # save the state of Gcf.figs
     init_figs = copy.copy(Gcf.figs)
 
@@ -51,18 +65,65 @@ def test_fig_close():
 
 
 @pytest.mark.backend('QtAgg', skip_on_importerror=True)
-def test_fig_signals(qt_core):
+@pytest.mark.parametrize("target, kwargs", [
+    (plt.show, {"block": True}),
+    (plt.pause, {"interval": 10})
+])
+def test_sigint(qt_core, platform_simulate_ctrl_c, target,
+                kwargs):
+    plt.figure()
+    def fire_signal():
+        platform_simulate_ctrl_c()
+
+    qt_core.QTimer.singleShot(100, fire_signal)
+    with pytest.raises(KeyboardInterrupt):
+        target(**kwargs)
+
+
+@pytest.mark.backend('QtAgg', skip_on_importerror=True)
+@pytest.mark.parametrize("target, kwargs", [
+    (plt.show, {"block": True}),
+    (plt.pause, {"interval": 10})
+])
+def test_other_signal_before_sigint(qt_core, platform_simulate_ctrl_c,
+                                    target, kwargs):
+    plt.figure()
+
+    sigcld_caught = False
+    def custom_sigpipe_handler(signum, frame):
+        nonlocal sigcld_caught
+        sigcld_caught = True
+    signal.signal(signal.SIGCLD, custom_sigpipe_handler)
+
+    def fire_other_signal():
+        os.kill(os.getpid(), signal.SIGCLD)
+
+    def fire_sigint():
+        platform_simulate_ctrl_c()
+
+    qt_core.QTimer.singleShot(50, fire_other_signal)
+    qt_core.QTimer.singleShot(100, fire_sigint)
+
+    with pytest.raises(KeyboardInterrupt):
+        target(**kwargs)
+
+    assert sigcld_caught
+
+
+@pytest.mark.backend('Qt5Agg')
+def test_fig_sigint_override(qt_core):
+    from matplotlib.backends.backend_qt5 import _BackendQT5
     # Create a figure
     plt.figure()
 
-    # Access signals
-    event_loop_signal = None
+    # Variable to access the handler from the inside of the event loop
+    event_loop_handler = None
 
     # Callback to fire during event loop: save SIGINT handler, then exit
     def fire_signal_and_quit():
         # Save event loop signal
-        nonlocal event_loop_signal
-        event_loop_signal = signal.getsignal(signal.SIGINT)
+        nonlocal event_loop_handler
+        event_loop_handler = signal.getsignal(signal.SIGINT)
 
         # Request event loop exit
         qt_core.QCoreApplication.exit()
@@ -71,26 +132,37 @@ def fire_signal_and_quit():
     qt_core.QTimer.singleShot(0, fire_signal_and_quit)
 
     # Save original SIGINT handler
-    original_signal = signal.getsignal(signal.SIGINT)
+    original_handler = signal.getsignal(signal.SIGINT)
 
     # Use our own SIGINT handler to be 100% sure this is working
-    def CustomHandler(signum, frame):
+    def custom_handler(signum, frame):
         pass
 
-    signal.signal(signal.SIGINT, CustomHandler)
+    signal.signal(signal.SIGINT, custom_handler)
 
     # mainloop() sets SIGINT, starts Qt event loop (which triggers timer and
     # exits) and then mainloop() resets SIGINT
     matplotlib.backends.backend_qt._BackendQT.mainloop()
 
-    # Assert: signal handler during loop execution is signal.SIG_DFL
-    assert event_loop_signal == signal.SIG_DFL
+    # Assert: signal handler during loop execution is changed
+    # (can't test equality with func)
+    assert event_loop_handler != custom_handler
 
     # Assert: current signal handler is the same as the one we set before
-    assert CustomHandler == signal.getsignal(signal.SIGINT)
+    assert signal.getsignal(signal.SIGINT) == custom_handler
+
+    # Repeat again to test that SIG_DFL and SIG_IGN will not be overridden
+    for custom_handler in (signal.SIG_DFL, signal.SIG_IGN):
+        qt_core.QTimer.singleShot(0, fire_signal_and_quit)
+        signal.signal(signal.SIGINT, custom_handler)
+
+        _BackendQT5.mainloop()
+
+        assert event_loop_handler == custom_handler
+        assert signal.getsignal(signal.SIGINT) == custom_handler
 
     # Reset SIGINT handler to what it was before the test
-    signal.signal(signal.SIGINT, original_signal)
+    signal.signal(signal.SIGINT, original_handler)
 
 
 @pytest.mark.parametrize(
