Add orjson_packer and orjson_unpacker to speed up serialization of session messages.

Author: fleming79

CHANGELOG.md

@@ -9,6 +9,7 @@
 ### Enhancements made
 
 - Support psutil for finding network addresses [#1033](https://github.com/jupyter/jupyter_client/pull/1033) ([@juliangilbey](https://github.com/juliangilbey))
+- Add new functions `orjson_packer` and `orjson_unpacker`.
 
 ### Bugs fixed
 

jupyter_client/session.py

@@ -32,6 +32,7 @@
 from traitlets import (
     Any,
     Bool,
+    Callable,
     CBytes,
     CUnicode,
     Dict,
@@ -124,15 +125,35 @@ def json_unpacker(s: str | bytes) -> t.Any:
     return json.loads(s)
 
 
+try:
+    import orjson  # type:ignore[import-not-found]
+except ModuleNotFoundError:
+    orjson = None
+    _default_packer_unpacker = "json", "json"
+    _default_pack_unpack = (json_packer, json_unpacker)
+else:
+
+    def orjson_packer(obj: t.Any) -> bytes:
+        """Convert a json object to a bytes using orjson with a fallback to json_packer."""
+        try:
+            return orjson.dumps(
+                obj, default=json_default, option=orjson.OPT_NAIVE_UTC | orjson.OPT_UTC_Z
+            )
+        except (TypeError, ValueError):
+            return json_packer(obj)
+
+    orjson_unpacker = orjson.loads
+    _default_packer_unpacker = "orjson", "orjson"
+    _default_pack_unpack = (orjson_packer, orjson_unpacker)
+
+
 def pickle_packer(o: t.Any) -> bytes:
     """Pack an object using the pickle module."""
     return pickle.dumps(squash_dates(o), PICKLE_PROTOCOL)
 
 
 pickle_unpacker = pickle.loads
 
-default_packer = json_packer
-default_unpacker = json_unpacker
 
 DELIM = b"<IDS|MSG>"
 # singleton dummy tracker, which will always report as done
@@ -350,48 +371,43 @@ class Session(Configurable):
         """,
     )
 
+    # serialization traits:
     packer = DottedObjectName(
-        "json",
+        _default_packer_unpacker[0],
         config=True,
         help="""The name of the packer for serializing messages.
             Should be one of 'json', 'pickle', or an import name
             for a custom callable serializer.""",
     )
-
-    @observe("packer")
-    def _packer_changed(self, change: t.Any) -> None:
-        new = change["new"]
-        if new.lower() == "json":
-            self.pack = json_packer
-            self.unpack = json_unpacker
-            self.unpacker = new
-        elif new.lower() == "pickle":
-            self.pack = pickle_packer
-            self.unpack = pickle_unpacker
-            self.unpacker = new
-        else:
-            self.pack = import_item(str(new))
-
     unpacker = DottedObjectName(
-        "json",
+        _default_packer_unpacker[1],
         config=True,
         help="""The name of the unpacker for unserializing messages.
         Only used with custom functions for `packer`.""",
     )
+    pack = Callable(_default_pack_unpack[0])  # the actual packer function
+    unpack = Callable(_default_pack_unpack[1])  # the actual unpacker function
 
-    @observe("unpacker")
-    def _unpacker_changed(self, change: t.Any) -> None:
+    @observe("packer", "unpacker")
+    def _packer_unpacker_changed(self, change: t.Any) -> None:
         new = change["new"]
-        if new.lower() == "json":
+        new_ = new.lower()
+        if new_ == "orjson" and orjson:
+            self.pack = orjson_packer
+            self.unpack = orjson_unpacker
+            self.packer = self.unpacker = new
+        elif new_ in ["json", "orjson"]:
             self.pack = json_packer
             self.unpack = json_unpacker
-            self.packer = new
-        elif new.lower() == "pickle":
+            self.packer = self.unpacker = new
+        elif new_ == "pickle":
             self.pack = pickle_packer
             self.unpack = pickle_unpacker
-            self.packer = new
+            self.packer = self.unpacker = new
         else:
-            self.unpack = import_item(str(new))
+            obj = import_item(str(new))
+            name = "pack" if change["name"] == "packer" else "unpack"
+            self.set_trait(name, obj)
 
     session = CUnicode("", config=True, help="""The UUID identifying this session.""")
 
@@ -416,8 +432,7 @@ def _session_changed(self, change: t.Any) -> None:
     metadata = Dict(
         {},
         config=True,
-        help="Metadata dictionary, which serves as the default top-level metadata dict for each "
-        "message.",
+        help="Metadata dictionary, which serves as the default top-level metadata dict for each message.",
     )
 
     # if 0, no adapting to do.
@@ -486,25 +501,6 @@ def _keyfile_changed(self, change: t.Any) -> None:
     # for protecting against sends from forks
     pid = Integer()
 
-    # serialization traits:
-
-    pack = Any(default_packer)  # the actual packer function
-
-    @observe("pack")
-    def _pack_changed(self, change: t.Any) -> None:
-        new = change["new"]
-        if not callable(new):
-            raise TypeError("packer must be callable, not %s" % type(new))
-
-    unpack = Any(default_unpacker)  # the actual packer function
-
-    @observe("unpack")
-    def _unpack_changed(self, change: t.Any) -> None:
-        # unpacker is not checked - it is assumed to be
-        new = change["new"]
-        if not callable(new):
-            raise TypeError("unpacker must be callable, not %s" % type(new))
-
     # thresholds:
     copy_threshold = Integer(
         2**16,
@@ -514,8 +510,7 @@ def _unpack_changed(self, change: t.Any) -> None:
     buffer_threshold = Integer(
         MAX_BYTES,
         config=True,
-        help="Threshold (in bytes) beyond which an object's buffer should be extracted to avoid "
-        "pickling.",
+        help="Threshold (in bytes) beyond which an object's buffer should be extracted to avoid pickling.",
     )
     item_threshold = Integer(
         MAX_ITEMS,
@@ -625,10 +620,7 @@ def _check_packers(self) -> None:
             unpacked = unpack(packed)
             assert unpacked == msg_list
         except Exception as e:
-            msg = (
-                f"unpacker '{self.unpacker}' could not handle output from packer"
-                f" '{self.packer}': {e}"
-            )
+            msg = f"unpacker '{self.unpacker}' could not handle output from packer '{self.packer}': {e}"
             raise ValueError(msg) from e
 
         # check datetime support

pyproject.toml

@@ -22,6 +22,7 @@ requires-python = ">=3.9"
 dependencies = [
     "importlib_metadata>=4.8.3;python_version<\"3.10\"",
     "jupyter_core>=4.12,!=5.0.*",
+    "orjson>=3.10.18; implementation_name == 'cpython'",
     "python-dateutil>=2.8.2",
     "pyzmq>=23.0",
     "tornado>=6.2",

tests/test_session.py

@@ -14,6 +14,7 @@
 import zmq
 from dateutil.tz import tzlocal
 from tornado import ioloop
+from traitlets import TraitError
 from zmq.eventloop.zmqstream import ZMQStream
 
 from jupyter_client import jsonutil
@@ -40,6 +41,14 @@ def session():
     return ss.Session()
 
 
+serializers = [
+    ("json", ss.json_packer, ss.json_unpacker),
+    ("pickle", ss.pickle_packer, ss.pickle_unpacker),
+]
+if ss.orjson:
+    serializers.append(("orjson", ss.orjson_packer, ss.orjson_unpacker))
+
+
 @pytest.mark.usefixtures("no_copy_threshold")
 class TestSession:
     def assertEqual(self, a, b):
@@ -63,7 +72,11 @@ def test_msg(self, session):
         self.assertEqual(msg["header"]["msg_type"], "execute")
         self.assertEqual(msg["msg_type"], "execute")
 
-    def test_serialize(self, session):
+    @pytest.mark.parametrize(["packer", "pack", "unpack"], serializers)
+    def test_serialize(self, session, packer, pack, unpack):
+        session.packer = packer
+        assert session.pack is pack
+        assert session.unpack is unpack
         msg = session.msg("execute", content=dict(a=10, b=1.1))
         msg_list = session.serialize(msg, ident=b"foo")
         ident, msg_list = session.feed_identities(msg_list)
@@ -233,16 +246,16 @@ async def test_send(self, session):
     def test_args(self, session):
         """initialization arguments for Session"""
         s = session
-        self.assertTrue(s.pack is ss.default_packer)
-        self.assertTrue(s.unpack is ss.default_unpacker)
+        assert s.pack is ss._default_pack_unpack[0]
+        assert s.unpack is ss._default_pack_unpack[1]
         self.assertEqual(s.username, os.environ.get("USER", "username"))
 
         s = ss.Session()
         self.assertEqual(s.username, os.environ.get("USER", "username"))
 
-        with pytest.raises(TypeError):
+        with pytest.raises(TraitError):
             ss.Session(pack="hi")
-        with pytest.raises(TypeError):
+        with pytest.raises(TraitError):
             ss.Session(unpack="hi")
         u = str(uuid.uuid4())
         s = ss.Session(username="carrot", session=u)
@@ -490,11 +503,6 @@ async def test_send_raw(self, session):
         B.close()
         ctx.term()
 
-    def test_set_packer(self, session):
-        s = session
-        s.packer = "json"
-        s.unpacker = "json"
-
     def test_clone(self, session):
         s = session
         s._add_digest("initial")
@@ -514,14 +522,43 @@ def test_squash_unicode():
     assert ss.squash_unicode("hi") == b"hi"
 
 
-def test_json_packer():
-    ss.json_packer(dict(a=1))
-    with pytest.raises(ValueError):
-        ss.json_packer(dict(a=ss.Session()))
-    ss.json_packer(dict(a=datetime(2021, 4, 1, 12, tzinfo=tzlocal())))
+@pytest.mark.parametrize(
+    ["description", "data"],
+    [
+        ("dict", [{"a": 1}, [{"a": 1}]]),
+        ("infinite", [math.inf, ["inf", None]]),
+        ("datetime", [datetime(2021, 4, 1, 12, tzinfo=tzlocal()), ["2021-04-01T12:00:00+11:00"]]),
+    ],
+)
+@pytest.mark.parametrize(["packer", "pack", "unpack"], serializers)
+def test_serialize_objects(packer, pack, unpack, description, data):
+    data_in, data_out_options = data
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
-        ss.json_packer(dict(a=math.inf))
+        value = pack(data_in)
+    unpacked = unpack(value)
+    if (description == "infinite") and (packer == "pickle"):
+        assert math.isinf(unpacked)
+        return
+    assert unpacked in data_out_options
+
+
+@pytest.mark.parametrize(["packer", "pack", "unpack"], serializers)
+def test_cannot_serialize(session, packer, pack, unpack):
+    data = {"a": session}
+    with pytest.raises((TypeError, ValueError)):
+        pack(data)
+
+
+@pytest.mark.parametrize("mode", ["packer", "unpacker"])
+@pytest.mark.parametrize(["packer", "pack", "unpack"], serializers)
+def test_packer_unpacker(session, packer, pack, unpack, mode):
+    s: ss.Session = session
+    s.set_trait(mode, packer)
+    assert s.pack is pack
+    assert s.unpack is unpack
+    mode_reverse = "unpacker" if mode == "packer" else "packer"
+    assert getattr(s, mode_reverse) == packer
 
 
 def test_message_cls():