Add msgpack support for message serialization in Session

Author: fleming79

CHANGELOG.md

@@ -9,7 +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`.
+- Add new functions `orjson_packer`, `orjson_unpacker`, `msgpack_packer` and `msgpack_unpacker`
 
 ### Bugs fixed
 

jupyter_client/session.py

@@ -12,6 +12,7 @@
 # Distributed under the terms of the Modified BSD License.
 from __future__ import annotations
 
+import functools
 import hashlib
 import hmac
 import json
@@ -132,15 +133,22 @@ def json_unpacker(s: str | bytes) -> t.Any:
     _default_packer_unpacker = "json", "json"
     _default_pack_unpack = (json_packer, json_unpacker)
 else:
-    import functools
-
     orjson_packer = functools.partial(
         orjson.dumps, default=json_default, option=orjson.OPT_NAIVE_UTC | orjson.OPT_UTC_Z
     )
     orjson_unpacker = orjson.loads
     _default_packer_unpacker = "orjson", "orjson"
     _default_pack_unpack = (orjson_packer, orjson_unpacker)
 
+try:
+    import msgpack  # type:ignore[import-not-found]
+
+except ModuleNotFoundError:
+    msgpack = None
+else:
+    msgpack_packer = functools.partial(msgpack.packb, default=json_default)
+    msgpack_unpacker = msgpack.unpackb
+
 
 def pickle_packer(o: t.Any) -> bytes:
     """Pack an object using the pickle module."""
@@ -331,7 +339,7 @@ class Session(Configurable):
 
     debug : bool
         whether to trigger extra debugging statements
-    packer/unpacker : str : 'json', 'pickle' or import_string
+    packer/unpacker : str : 'orjson', 'json', 'pickle', 'msgpack' or import_string
         importstrings for methods to serialize message parts.  If just
         'json' or 'pickle', predefined JSON and pickle packers will be used.
         Otherwise, the entire importstring must be used.
@@ -399,6 +407,10 @@ def _packer_unpacker_changed(self, change: t.Any) -> None:
             self.pack = pickle_packer
             self.unpack = pickle_unpacker
             self.packer = self.unpacker = new
+        elif new_ == "msgpack":
+            self.pack = msgpack_packer
+            self.unpack = msgpack_unpacker
+            self.packer = self.unpacker = new
         else:
             obj = import_item(str(new))
             name = "pack" if change["name"] == "packer" else "unpack"
@@ -523,7 +535,7 @@ def __init__(self, **kwargs: t.Any) -> None:
 
         debug : bool
             whether to trigger extra debugging statements
-        packer/unpacker : str : 'json', 'pickle' or import_string
+        packer/unpacker : str : 'orjson', 'json', 'pickle', 'msgpack' or import_string
             importstrings for methods to serialize message parts.  If just
             'json' or 'pickle', predefined JSON and pickle packers will be used.
             Otherwise, the entire importstring must be used.

tests/test_session.py

@@ -47,6 +47,8 @@ def session():
 ]
 if ss.orjson:
     serializers.append(("orjson", ss.orjson_packer, ss.orjson_unpacker))
+if ss.msgpack:
+    serializers.append(("msgpack", ss.msgpack_packer, ss.msgpack_unpacker))
 
 
 @pytest.mark.usefixtures("no_copy_threshold")
@@ -537,7 +539,7 @@ def test_serialize_objects(packer, pack, unpack, description, data):
         warnings.simplefilter("ignore")
         value = pack(data_in)
     unpacked = unpack(value)
-    if (description == "infinite") and (packer == "pickle"):
+    if (description == "infinite") and (packer in ["pickle", "msgpack"]):
         assert math.isinf(unpacked)
         return
     assert unpacked in data_out_options
@@ -552,7 +554,7 @@ def test_cannot_serialize(session, packer, pack, unpack):
 
 @pytest.mark.parametrize("mode", ["packer", "unpacker"])
 @pytest.mark.parametrize(["packer", "pack", "unpack"], serializers)
-def test_packer_unpacker(session, packer, pack, unpack, mode):
+def test_pack_unpack(session, packer, pack, unpack, mode):
     s: ss.Session = session
     s.set_trait(mode, packer)
     assert s.pack is pack