Add asyncio_event_loop (currently only set for asyncio backend) to the kernel and drop old integration code.

Author: fleming79

ipykernel/eventloops.py

@@ -6,7 +6,6 @@
 import os
 import platform
 import sys
-from functools import partial
 
 import zmq
 from packaging.version import Version as V
@@ -400,63 +399,11 @@ def loop_cocoa_exit(kernel):
 
 @register_integration("asyncio")
 def loop_asyncio(kernel):
-    """Start a kernel with asyncio event loop support."""
-    import asyncio
+    """Verify the asyncio event loop is supported."""
 
-    loop = asyncio.get_event_loop()
-    # loop is already running (e.g. tornado 5), nothing left to do
-    if loop.is_running():
-        return
-
-    if loop.is_closed():
-        # main loop is closed, create a new one
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-    loop._should_close = False  # type:ignore[attr-defined]
-
-    # pause eventloop when there's an event on a zmq socket
-    def process_stream_events(socket):
-        """fall back to main loop when there's a socket event"""
-        loop.stop()
-
-    notifier = partial(process_stream_events, kernel.shell_socket)
-    loop.add_reader(kernel.shell_socket.getsockopt(zmq.FD), notifier)
-    loop.call_soon(notifier)
-
-    while True:
-        error = None
-        try:
-            loop.run_forever()
-        except KeyboardInterrupt:
-            continue
-        except Exception as e:
-            error = e
-        if loop._should_close:  # type:ignore[attr-defined]
-            loop.close()
-        if error is not None:
-            raise error
-        break
-
-
-@loop_asyncio.exit
-def loop_asyncio_exit(kernel):
-    """Exit hook for asyncio"""
-    import asyncio
-
-    loop = asyncio.get_event_loop()
-
-    async def close_loop():
-        if hasattr(loop, "shutdown_asyncgens"):
-            yield loop.shutdown_asyncgens()
-        loop._should_close = True  # type:ignore[attr-defined]
-        loop.stop()
-
-    if loop.is_running():
-        close_loop()
-
-    elif not loop.is_closed():
-        loop.run_until_complete(close_loop)  # type:ignore[arg-type]
-        loop.close()
+    if not kernel.asyncio_event_loop or not kernel.asyncio_event_loop.is_running():
+        msg = "The asyncio event loop is not running so is not supported."
+        raise RuntimeError(msg)
 
 
 def set_qt_api_env_from_gui(gui):

ipykernel/kernelbase.py

@@ -4,6 +4,8 @@
 # Distributed under the terms of the Modified BSD License.
 from __future__ import annotations
 
+import asyncio
+import contextlib
 import inspect
 import itertools
 import logging
@@ -128,6 +130,7 @@ class Kernel(SingletonConfigurable):
 
     _send_exec_request: Dict[dict[zmq_anyio.Socket, MemoryObjectSendStream]] = Dict()
     _main_subshell_ready = Instance(Event, ())
+    asyncio_event_loop = Instance(asyncio.AbstractEventLoop, allow_none=True, read_only=True)  # type:ignore[call-overload]
 
     log: logging.Logger = Instance(logging.Logger, allow_none=True)  # type:ignore[assignment]
 
@@ -441,10 +444,13 @@ async def shell_main(self, subshell_id: str | None):
                 tg.start_soon(self._execute_request_handler, receive_stream)
                 if subshell_id is None:
                     # Main subshell.
+                    with contextlib.suppress(RuntimeError):
+                        self.set_trait("asyncio_event_loop", asyncio.get_running_loop())
                     self._main_subshell_ready.set()
                     await to_thread.run_sync(self.shell_stop.wait)
                     tg.cancel_scope.cancel()
             self._send_exec_request.pop(socket, None)
+            self.set_trait("asyncio_event_loop", None)
             await send_stream.aclose()
 
     async def _execute_request_handler(self, receive_stream: MemoryObjectReceiveStream):
@@ -821,7 +827,7 @@ async def execute_request(self, socket, ident, parent):
             self._aborted_time = time.monotonic()
             self.log.info("Aborting queue")
 
-    def do_execute(
+    async def do_execute(
         self,
         code,
         silent,

tests/test_eventloop.py

@@ -1,6 +1,5 @@
 """Test eventloop integration"""
 
-import asyncio
 import os
 import sys
 import threading
@@ -80,18 +79,12 @@ def do_thing():
     t.join(1)
 
 
-@windows_skip
-@pytest.mark.parametrize("anyio_backend", ["asyncio"])
-async def test_asyncio_loop(kernel):
-    okay = asyncio.Event()
-
-    def do_thing():
-        okay.set()
-
-    loop = asyncio.get_event_loop()
-    loop.call_soon(do_thing)
-    loop_asyncio(kernel)
-    await asyncio.wait_for(okay.wait(), 1)
+async def test_asyncio_loop(kernel, anyio_backend):
+    if anyio_backend == "asyncio":
+        loop_asyncio(kernel)
+    else:
+        with pytest.raises(RuntimeError):
+            loop_asyncio(kernel)
 
 
 @windows_skip