Fix matplotlib use of GUI toolkit event loops

Author: ianthomas23

ipykernel/eventloops.py

@@ -40,8 +40,9 @@ def register_integration(*toolkitnames):
     You can provide alternative names for the same toolkit.
 
     The decorated function should take a single argument, the IPython kernel
-    instance, arrange for the event loop to call ``kernel.do_one_iteration()``
-    at least every ``kernel._poll_interval`` seconds, and start the event loop.
+    instance, arrange for the event loop to yield the asyncio loop when a
+    message is received by the main shell zmq stream or at least every
+    ``kernel._poll_interval`` seconds, and start the event loop.
 
     :mod:`ipykernel.eventloops` provides and registers such functions
     for a few common event loops.
@@ -68,6 +69,15 @@ def exit_decorator(exit_func):
     return decorator
 
 
+def get_shell_stream(kernel):
+    # Return the zmq stream that receives messages for the main shell.
+    if kernel._supports_kernel_subshells:
+        manager = kernel.shell_channel_thread.manager
+        socket_pair = manager.get_shell_channel_to_subshell_pair(None)
+        return socket_pair.to_stream
+    return kernel.shell_stream
+
+
 def _notify_stream_qt(kernel):
     import operator
     from functools import lru_cache
@@ -87,17 +97,20 @@ def exit_loop():
         kernel._qt_notifier.setEnabled(False)
         kernel.app.qt_event_loop.quit()
 
-    def process_stream_events():
+    def process_stream_events_wrap(shell_stream):
         """fall back to main loop when there's a socket event"""
         # call flush to ensure that the stream doesn't lose events
         # due to our consuming of the edge-triggered FD
         # flush returns the number of events consumed.
         # if there were any, wake it up
-        if kernel.shell_stream.flush(limit=1):
+        if shell_stream.flush(limit=1):
             exit_loop()
 
+    shell_stream = get_shell_stream(kernel)
+    process_stream_events = partial(process_stream_events_wrap, shell_stream)
+
     if not hasattr(kernel, "_qt_notifier"):
-        fd = kernel.shell_stream.getsockopt(zmq.FD)
+        fd = shell_stream.getsockopt(zmq.FD)
         kernel._qt_notifier = QtCore.QSocketNotifier(
             fd, enum_helper("QtCore.QSocketNotifier.Type").Read, kernel.app.qt_event_loop
         )
@@ -177,9 +190,11 @@ def loop_wx(kernel):
     # Wx uses milliseconds
     poll_interval = int(1000 * kernel._poll_interval)
 
-    def wake():
+    shell_stream = get_shell_stream(kernel)
+
+    def wake(shell_stream):
         """wake from wx"""
-        if kernel.shell_stream.flush(limit=1):
+        if shell_stream.flush(limit=1):
             kernel.app.ExitMainLoop()
             return
 
@@ -201,7 +216,7 @@ def on_timer(self, event):
     # wx.Timer to defer back to the tornado event loop.
     class IPWxApp(wx.App):  # type:ignore[misc]
         def OnInit(self):
-            self.frame = TimerFrame(wake)
+            self.frame = TimerFrame(partial(wake, shell_stream))
             self.frame.Show(False)
             return True
 
@@ -248,14 +263,14 @@ def __init__(self, app):
 
         def exit_loop():
             """fall back to main loop"""
-            app.tk.deletefilehandler(kernel.shell_stream.getsockopt(zmq.FD))
+            app.tk.deletefilehandler(shell_stream.getsockopt(zmq.FD))
             app.quit()
             app.destroy()
             del kernel.app_wrapper
 
-        def process_stream_events(*a, **kw):
+        def process_stream_events_wrap(shell_stream, *a, **kw):
             """fall back to main loop when there's a socket event"""
-            if kernel.shell_stream.flush(limit=1):
+            if shell_stream.flush(limit=1):
                 exit_loop()
 
         # allow for scheduling exits from the loop in case a timeout needs to
@@ -268,9 +283,10 @@ def _schedule_exit(delay):
 
         # For Tkinter, we create a Tk object and call its withdraw method.
         kernel.app_wrapper = BasicAppWrapper(app)
-        app.tk.createfilehandler(
-            kernel.shell_stream.getsockopt(zmq.FD), READABLE, process_stream_events
-        )
+        shell_stream = get_shell_stream(kernel)
+        process_stream_events = partial(process_stream_events_wrap, shell_stream)
+
+        app.tk.createfilehandler(shell_stream.getsockopt(zmq.FD), READABLE, process_stream_events)
         # schedule initial call after start
         app.after(0, process_stream_events)
 
@@ -283,15 +299,19 @@ def _schedule_exit(delay):
 
         nest_asyncio.apply()
 
-        doi = kernel.do_one_iteration
         # Tk uses milliseconds
         poll_interval = int(1000 * kernel._poll_interval)
 
+        shell_stream = get_shell_stream(kernel)
+
         class TimedAppWrapper:
-            def __init__(self, app, func):
+            def __init__(self, app, shell_stream):
                 self.app = app
+                self.shell_stream = shell_stream
                 self.app.withdraw()
-                self.func = func
+
+            async def func(self):
+                self.shell_stream.flush(limit=1)
 
             def on_timer(self):
                 loop = asyncio.get_event_loop()
@@ -305,16 +325,18 @@ def start(self):
                 self.on_timer()  # Call it once to get things going.
                 self.app.mainloop()
 
-        kernel.app_wrapper = TimedAppWrapper(app, doi)
+        kernel.app_wrapper = TimedAppWrapper(app, shell_stream)
         kernel.app_wrapper.start()
 
 
 @loop_tk.exit
 def loop_tk_exit(kernel):
     """Exit the tk loop."""
     try:
+        kernel.app_wrapper.app.quit()
         kernel.app_wrapper.app.destroy()
         del kernel.app_wrapper
+        kernel.eventloop = None
     except (RuntimeError, AttributeError):
         pass
 
@@ -359,6 +381,7 @@ def loop_cocoa(kernel):
     from ._eventloop_macos import mainloop, stop
 
     real_excepthook = sys.excepthook
+    shell_stream = get_shell_stream(kernel)
 
     def handle_int(etype, value, tb):
         """don't let KeyboardInterrupts look like crashes"""
@@ -377,7 +400,7 @@ def handle_int(etype, value, tb):
                 # don't let interrupts during mainloop invoke crash_handler:
                 sys.excepthook = handle_int
                 mainloop(kernel._poll_interval)
-                if kernel.shell_stream.flush(limit=1):
+                if shell_stream.flush(limit=1):
                     # events to process, return control to kernel
                     return
             except BaseException:
@@ -415,13 +438,14 @@ def loop_asyncio(kernel):
     loop._should_close = False  # type:ignore[attr-defined]
 
     # pause eventloop when there's an event on a zmq socket
-    def process_stream_events(stream):
+    def process_stream_events(shell_stream):
         """fall back to main loop when there's a socket event"""
-        if stream.flush(limit=1):
+        if shell_stream.flush(limit=1):
             loop.stop()
 
-    notifier = partial(process_stream_events, kernel.shell_stream)
-    loop.add_reader(kernel.shell_stream.getsockopt(zmq.FD), notifier)
+    shell_stream = get_shell_stream(kernel)
+    notifier = partial(process_stream_events, shell_stream)
+    loop.add_reader(shell_stream.getsockopt(zmq.FD), notifier)
     loop.call_soon(notifier)
 
     while True:

ipykernel/kernelbase.py

@@ -6,7 +6,6 @@
 
 import asyncio
 import inspect
-import itertools
 import logging
 import os
 import socket
@@ -42,7 +41,6 @@
 from IPython.core.error import StdinNotImplementedError
 from jupyter_client.session import Session
 from tornado import ioloop
-from tornado.queues import Queue, QueueEmpty
 from traitlets.config.configurable import SingletonConfigurable
 from traitlets.traitlets import (
     Any,
@@ -507,11 +505,6 @@ async def advance_eventloop():
             if self.eventloop is not eventloop:
                 self.log.info("exiting eventloop %s", eventloop)
                 return
-            if self.msg_queue.qsize():
-                self.log.debug("Delaying eventloop due to waiting messages")
-                # still messages to process, make the eventloop wait
-                schedule_next()
-                return
             self.log.debug("Advancing eventloop %s", eventloop)
             try:
                 eventloop(self)
@@ -530,98 +523,18 @@ def schedule_next():
             # already consumed from the queue by process_one and the queue is
             # technically empty.
             self.log.debug("Scheduling eventloop advance")
-            self.io_loop.call_later(0.001, partial(self.schedule_dispatch, advance_eventloop))
+            self.io_loop.call_later(0.001, advance_eventloop)
 
         # begin polling the eventloop
         schedule_next()
 
-    async def do_one_iteration(self):
-        """Process a single shell message
-
-        Any pending control messages will be flushed as well
-
-        .. versionchanged:: 5
-            This is now a coroutine
-        """
-        # flush messages off of shell stream into the message queue
-        if self.shell_stream and not self._supports_kernel_subshells:
-            self.shell_stream.flush()
-        # process at most one shell message per iteration
-        await self.process_one(wait=False)
-
-    async def process_one(self, wait=True):
-        """Process one request
-
-        Returns None if no message was handled.
-        """
-        if wait:
-            t, dispatch, args = await self.msg_queue.get()
-        else:
-            try:
-                t, dispatch, args = self.msg_queue.get_nowait()
-            except (asyncio.QueueEmpty, QueueEmpty):
-                return
-
-        if self.control_thread is None and self.control_stream is not None:
-            # If there isn't a separate control thread then this main thread handles both shell
-            # and control messages. Before processing a shell message we need to flush all control
-            # messages and allow them all to be processed.
-            await asyncio.sleep(0)
-            self.control_stream.flush()
-
-            socket = self.control_stream.socket
-            while socket.poll(1):
-                await asyncio.sleep(0)
-                self.control_stream.flush()
-
-        await dispatch(*args)
-
-    async def dispatch_queue(self):
-        """Coroutine to preserve order of message handling
-
-        Ensures that only one message is processing at a time,
-        even when the handler is async
-        """
-
-        while True:
-            try:
-                await self.process_one()
-            except Exception:
-                self.log.exception("Error in message handler")
-
-    _message_counter = Any(
-        help="""Monotonic counter of messages
-        """,
-    )
-
-    @default("_message_counter")
-    def _message_counter_default(self):
-        return itertools.count()
-
-    def schedule_dispatch(self, dispatch, *args):
-        """schedule a message for dispatch"""
-        idx = next(self._message_counter)
-
-        self.msg_queue.put_nowait(
-            (
-                idx,
-                dispatch,
-                args,
-            )
-        )
-        # ensure the eventloop wakes up
-        self.io_loop.add_callback(lambda: None)
-
     async def _create_control_lock(self):
         # This can be removed when minimum python increases to 3.10
         self._control_lock = asyncio.Lock()
 
     def start(self):
         """register dispatchers for streams"""
         self.io_loop = ioloop.IOLoop.current()
-        self.msg_queue: Queue[t.Any] = Queue()
-        if not self.shell_channel_thread:
-            self.io_loop.add_callback(self.dispatch_queue)
 
         if self.control_stream:
             self.control_stream.on_recv(self.dispatch_control, copy=False)
@@ -640,10 +553,7 @@ def start(self):
                 self.shell_stream.on_recv(self.shell_channel_thread_main, copy=False)
             else:
                 self.shell_stream.on_recv(
-                    partial(
-                        self.schedule_dispatch,
-                        self.dispatch_shell,
-                    ),
+                    partial(self.shell_main, None),
                     copy=False,
                 )
 
@@ -1391,23 +1301,18 @@ async def stop_aborting():
             self.log.info("Finishing abort")
             self._aborting = False
 
-        # put the stop-aborting event on the message queue
-        # so that all messages already waiting in the queue are aborted
-        # before we reset the flag
-        schedule_stop_aborting = partial(self.schedule_dispatch, stop_aborting)
-
         if self.stop_on_error_timeout:
             # if we have a delay, give messages this long to arrive on the queue
             # before we stop aborting requests
-            self.io_loop.call_later(self.stop_on_error_timeout, schedule_stop_aborting)
+            self.io_loop.call_later(self.stop_on_error_timeout, stop_aborting)
             # If we have an eventloop, it may interfere with the call_later above.
             # If the loop has a _schedule_exit method, we call that so the loop exits
             # after stop_on_error_timeout, returning to the main io_loop and letting
             # the call_later fire.
             if self.eventloop is not None and hasattr(self.eventloop, "_schedule_exit"):
                 self.eventloop._schedule_exit(self.stop_on_error_timeout + 0.01)
         else:
-            schedule_stop_aborting()
+            self.io_loop.add_callback(stop_aborting)
 
     def _send_abort_reply(self, stream, msg, idents):
         """Send a reply to an aborted request"""

tests/test_ipkernel_direct.py

@@ -165,29 +165,37 @@ def test_dispatch_debugpy(ipkernel: IPythonKernel) -> None:
 async def test_start(ipkernel: IPythonKernel) -> None:
     shell_future: asyncio.Future = asyncio.Future()
 
-    async def fake_dispatch_queue():
-        shell_future.set_result(None)
+    def fake_publish_status(status, channel):
+        if status == "starting" and channel == "shell":
+            shell_future.set_result(None)
 
-    ipkernel.dispatch_queue = fake_dispatch_queue  # type:ignore
+    ipkernel._publish_status = fake_publish_status # type:ignore
     ipkernel.start()
-    ipkernel.debugpy_stream = None
-    ipkernel.start()
-    await ipkernel.process_one(False)
+
     await shell_future
 
+    shell_stream = ipkernel.shell_stream
+    assert shell_stream is not None
+    assert not shell_stream.closed()
+
 
 async def test_start_no_debugpy(ipkernel: IPythonKernel) -> None:
     shell_future: asyncio.Future = asyncio.Future()
 
-    async def fake_dispatch_queue():
-        shell_future.set_result(None)
+    def fake_publish_status(status, channel):
+        if status == "starting" and channel == "shell":
+            shell_future.set_result(None)
 
-    ipkernel.dispatch_queue = fake_dispatch_queue  # type:ignore
+    ipkernel._publish_status = fake_publish_status # type:ignore
     ipkernel.debugpy_stream = None
     ipkernel.start()
 
     await shell_future
 
+    shell_stream = ipkernel.shell_stream
+    assert shell_stream is not None
+    assert not shell_stream.closed()
+
 
 def test_create_comm():
     assert isinstance(_create_comm(), BaseComm)