Make kernel independent of zmq (#45)

Author: davidbrochart

src/akernel/akernel.py

@@ -9,7 +9,6 @@
 from .connect import connect_channel
 from .kernel import Kernel
 from .kernelspec import write_kernelspec
-from .message import receive_message, send_message
 
 
 cli = typer.Typer()
@@ -106,34 +105,34 @@ async def start(self) -> None:
 
     async def to_shell(self) -> None:
         while True:
-            msg = await receive_message(self.shell_channel)
+            msg = await self.shell_channel.arecv_multipart().wait()
             await self._to_shell_send_stream.send(msg)
 
     async def from_shell(self) -> None:
         async for msg in self._from_shell_receive_stream:
-            await send_message(msg, self.shell_channel)
+            await self.shell_channel.asend_multipart(msg, copy=True).wait()
 
     async def to_control(self) -> None:
         while True:
-            msg = await receive_message(self.control_channel)
+            msg = await self.control_channel.arecv_multipart().wait()
             await self._to_control_send_stream.send(msg)
 
     async def from_control(self) -> None:
         async for msg in self._from_control_receive_stream:
-            await send_message(msg, self.control_channel)
+            await self.control_channel.asend_multipart(msg, copy=True).wait()
 
     async def to_stdin(self) -> None:
         while True:
-            msg = await receive_message(self.stdin_channel)
+            msg = await self.stdin_channel.arecv_multipart().wait()
             await self._to_stdin_send_stream.send(msg)
 
     async def from_stdin(self) -> None:
         async for msg in self._from_stdin_receive_stream:
-            await send_message(msg, self.stdin_channel)
+            await self.stdin_channel.asend_multipart(msg, copy=True).wait()
 
     async def from_iopub(self) -> None:
         async for msg in self._from_iopub_receive_stream:
-            await send_message(msg, self.iopub_channel)
+            await self.iopub_channel.asend_multipart(msg, copy=True).wait()
 
 
 if __name__ == "__main__":

src/akernel/kernel.py

@@ -14,7 +14,6 @@
 from akernel.display import display
 import akernel.IPython
 from akernel.IPython import core
-from .connect import connect_channel
 from .message import create_message, feed_identities, deserialize, serialize
 from .execution import pre_execute, cache_execution
 from .traceback import get_traceback

src/akernel/message.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
+import json
 import uuid
 import hmac
 import hashlib
 from datetime import datetime, timezone
 from typing import Any, cast
 
-from zmq.utils import jsonapi
-from zmq_anyio import Socket
 from dateutil.parser import parse as dateutil_parse  # type: ignore
 
 
@@ -97,24 +96,30 @@ def serialize(msg: dict[str, Any], key: str) -> list[bytes]:
     return to_send
 
 
-async def receive_message(sock: Socket) -> tuple[list[bytes], dict[str, Any]] | None:
-    return await sock.arecv_multipart().wait()
-    return None
+def dumps(o: Any, **kwargs) -> bytes:
+    """Serialize object to JSON bytes (utf-8).
 
+    Keyword arguments are passed along to :py:func:`json.dumps`.
+    """
+    return json.dumps(o, **kwargs).encode("utf8")
 
-async def send_message(
-    msg: dict[str, Any],
-    sock: Socket,
-) -> None:
-    await sock.asend_multipart(msg, copy=True).wait()
+
+def loads(s: bytes | str, **kwargs) -> dict | list | str | int | float:
+    """Load object from JSON bytes (utf-8).
+
+    Keyword arguments are passed along to :py:func:`json.loads`.
+    """
+    if isinstance(s, bytes):
+        s = s.decode("utf8")
+    return json.loads(s, **kwargs)
 
 
 def pack(obj: dict[str, Any]) -> bytes:
-    return jsonapi.dumps(obj)
+    return dumps(obj)
 
 
 def unpack(s: bytes) -> dict[str, Any]:
-    return cast(dict[str, Any], jsonapi.loads(s))
+    return cast(dict[str, Any], loads(s))
 
 
 def sign(msg_list: list[bytes], key: str) -> bytes: