Use zmq-anyio v0.3.0

Author: davidbrochart

ipykernel/debugger.py

@@ -241,7 +241,7 @@ async def _send_request(self, msg):
         self.log.debug("DEBUGPYCLIENT:")
         self.log.debug(self.routing_id)
         self.log.debug(buf)
-        await self.debugpy_socket.asend_multipart((self.routing_id, buf))
+        await self.debugpy_socket.asend_multipart((self.routing_id, buf)).wait()
 
     async def _wait_for_response(self):
         # Since events are never pushed to the message_queue
@@ -437,7 +437,7 @@ async def start(self):
                 (self.shell_socket.getsockopt(ROUTING_ID)),
             )
 
-            msg = await self.shell_socket.arecv_multipart()
+            msg = await self.shell_socket.arecv_multipart().wait()
             ident, msg = self.session.feed_identities(msg, copy=True)
             try:
                 msg = self.session.deserialize(msg, content=True, copy=True)

ipykernel/inprocess/session.py

@@ -3,7 +3,7 @@
 
 class Session(_Session):
     async def recv(self, socket, copy=True):
-        return await socket.arecv_multipart()
+        return await socket.arecv_multipart().wait()
 
     def send(
         self,

ipykernel/iostream.py

@@ -143,7 +143,7 @@ async def _handle_event(self):
         async with pipe_in:
             try:
                 while True:
-                    await pipe_in.arecv()
+                    await pipe_in.arecv().wait()
                     # freeze event count so new writes don't extend the queue
                     # while we are processing
                     n_events = len(self._events)
@@ -189,7 +189,7 @@ async def _handle_pipe_msgs(self):
 
     async def _handle_pipe_msg(self, msg=None):
         """handle a pipe message from a subprocess"""
-        msg = msg or await self._pipe_in1.arecv_multipart()
+        msg = msg or await self._pipe_in1.arecv_multipart().wait()
         if not self._pipe_flag or not self._is_main_process():
             return
         if msg[0] != self._pipe_uuid:

ipykernel/ipkernel.py

@@ -236,7 +236,7 @@ async def receive_debugpy_message(self, msg=None):
 
         if msg is None:
             assert self.debugpy_socket is not None
-            msg = await self.debugpy_socket.arecv_multipart()
+            msg = await self.debugpy_socket.arecv_multipart().wait()
         # The first frame is the socket id, we can drop it
         frame = msg[1].decode("utf-8")
         self.log.debug("Debugpy received: %s", frame)

ipykernel/kernelbase.py

@@ -269,7 +269,7 @@ async def process_control_message(self, msg=None):
         assert self.session is not None
         assert self.control_thread is None or threading.current_thread() == self.control_thread
 
-        msg = msg or await self.control_socket.arecv_multipart()
+        msg = msg or await self.control_socket.arecv_multipart().wait()
         idents, msg = self.session.feed_identities(msg)
         try:
             msg = self.session.deserialize(msg, content=True)
@@ -369,7 +369,7 @@ async def shell_channel_thread_main(self):
         async with self.shell_socket, create_task_group() as tg:
             try:
                 while True:
-                    msg = await self.shell_socket.arecv_multipart(copy=False)
+                    msg = await self.shell_socket.arecv_multipart(copy=False).wait()
                     # deserialize only the header to get subshell_id
                     # Keep original message to send to subshell_id unmodified.
                     _, msg2 = self.session.feed_identities(msg, copy=False)
@@ -384,7 +384,7 @@ async def shell_channel_thread_main(self):
                         assert socket is not None
                         if not socket.started.is_set():
                             await tg.start(socket.start)
-                        await socket.asend_multipart(msg, copy=False)
+                        socket.asend_multipart(msg, copy=False)
                     except Exception:
                         self.log.error("Invalid message", exc_info=True)  # noqa: G201
             except BaseException:
@@ -444,8 +444,8 @@ async def process_shell_message(self, msg=None, socket=None):
             assert socket is None
             socket = self.shell_socket
 
-        no_msg = msg is None if self._is_test else not await socket.apoll(0)
-        msg = msg or await socket.arecv_multipart(copy=False)
+        no_msg = msg is None if self._is_test else not await socket.apoll(0).wait()
+        msg = msg or await socket.arecv_multipart(copy=False).wait()
 
         received_time = time.monotonic()
         copy = not isinstance(msg[0], zmq.Message)
@@ -499,7 +499,7 @@ async def process_shell_message(self, msg=None, socket=None):
             try:
                 result = handler(socket, idents, msg)
                 if inspect.isawaitable(result):
-                    result = await result
+                    await result
             except Exception:
                 self.log.error("Exception in message handler:", exc_info=True)  # noqa: G201
             except KeyboardInterrupt:
@@ -1090,8 +1090,8 @@ async def create_subshell_request(self, socket, ident, parent) -> None:
         other_socket = await self.shell_channel_thread.manager.get_control_other_socket(
             self.control_thread
         )
-        await other_socket.asend_json({"type": "create"})
-        reply = await other_socket.arecv_json()
+        await other_socket.asend_json({"type": "create"}).wait()
+        reply = await other_socket.arecv_json().wait()
 
         self.session.send(socket, "create_subshell_reply", reply, parent, ident)
 
@@ -1114,8 +1114,8 @@ async def delete_subshell_request(self, socket, ident, parent) -> None:
         other_socket = await self.shell_channel_thread.manager.get_control_other_socket(
             self.control_thread
         )
-        await other_socket.asend_json({"type": "delete", "subshell_id": subshell_id})
-        reply = await other_socket.arecv_json()
+        await other_socket.asend_json({"type": "delete", "subshell_id": subshell_id}).wait()
+        reply = await other_socket.arecv_json().wait()
 
         self.session.send(socket, "delete_subshell_reply", reply, parent, ident)
 
@@ -1131,8 +1131,8 @@ async def list_subshell_request(self, socket, ident, parent) -> None:
         other_socket = await self.shell_channel_thread.manager.get_control_other_socket(
             self.control_thread
         )
-        await other_socket.asend_json({"type": "list"})
-        reply = await other_socket.arecv_json()
+        await other_socket.asend_json({"type": "list"}).wait()
+        reply = await other_socket.arecv_json().wait()
 
         self.session.send(socket, "list_subshell_reply", reply, parent, ident)
 

ipykernel/subshell_manager.py

@@ -158,9 +158,9 @@ async def listen_from_control(self, subshell_task: t.Any) -> None:
         socket = self._control_shell_channel_socket
         async with socket:
             while True:
-                request = await socket.arecv_json()
+                request = await socket.arecv_json().wait()
                 reply = await self._process_control_request(request, subshell_task)
-                await socket.asend_json(reply)
+                await socket.asend_json(reply).wait()
 
     async def listen_from_subshells(self) -> None:
         """Listen for reply messages on inproc sockets of all subshells and resend
@@ -270,9 +270,9 @@ async def _listen_for_subshell_reply(
         await shell_channel_socket.started.wait()
         try:
             while True:
-                msg = await shell_channel_socket.arecv_multipart(copy=False)
+                msg = await shell_channel_socket.arecv_multipart(copy=False).wait()
                 with self._lock_shell_socket:
-                    await self._shell_socket.asend_multipart(msg)
+                    await self._shell_socket.asend_multipart(msg).wait()
         except BaseException:
             if not self._is_subshell(subshell_id):
                 # Subshell no longer exists so exit gracefully

pyproject.toml

@@ -34,7 +34,7 @@ dependencies = [
     "psutil>=5.7",
     "packaging>=22",
     "anyio>=4.8.0,<5.0.0",
-    "zmq-anyio >=0.2.5",
+    "zmq-anyio >=0.3.0",
 ]
 
 [project.urls]

tests/test_io.py

@@ -221,8 +221,8 @@ async def test_echo_watch(ctx):
         print(f"{p.stdout=}")
         print(f"{p.stderr}=", file=sys.stderr)
         assert p.returncode == 0
-        while await s.apoll(timeout=100):
-            msg = await s.arecv_multipart()
+        while await s.apoll(timeout=100).wait():
+            msg = await s.arecv_multipart().wait()
             ident, msg = session.feed_identities(msg, copy=True)
             msg = session.deserialize(msg, content=True, copy=True)
             assert msg is not None  # for type narrowing

tests/test_kernel.py

@@ -83,8 +83,6 @@ def thread_target():
             msg = kc.get_iopub_msg(timeout=interval * 2)
             if msg["msg_type"] != "stream":
                 continue
-            print(f"{thread_msg_id=}")
-            print(f"{msg=}")
             content = msg["content"]
             assert content["name"] == "stdout"
             assert content["text"] == str(received)
@@ -123,8 +121,6 @@ def parent_target():
             msg = kc.get_iopub_msg(timeout=interval * 2)
             if msg["msg_type"] != "stream":
                 continue
-            print(f"{thread_msg_id=}")
-            print(f"{msg=}")
             content = msg["content"]
             assert content["name"] == "stdout"
             assert content["text"] == str(received)
@@ -157,8 +153,6 @@ async def async_task():
             msg = kc.get_iopub_msg(timeout=interval * 2)
             if msg["msg_type"] != "stream":
                 continue
-            print(f"{thread_msg_id=}")
-            print(f"{msg=}")
             content = msg["content"]
             assert content["name"] == "stdout"
             assert content["text"] == str(received)