serde.msgpack

  1"""
  2Serialize and Deserialize in MsgPack format. This module depends on
  3[msgpack](https://pypi.org/project/msgpack/) package.
  4"""
  5
  6from typing import Any, Optional, overload
  7
  8import msgpack
  9
 10from .compat import T
 11from .compat import SerdeError
 12from .de import Deserializer, from_dict, from_tuple
 13from .numpy import encode_numpy
 14from .se import Serializer, to_dict, to_tuple
 15
 16__all__ = ["from_msgpack", "to_msgpack"]
 17
 18
 19class MsgPackSerializer(Serializer[bytes]):
 20    @classmethod
 21    def serialize(
 22        cls, obj: Any, use_bin_type: bool = True, ext_type_code: Optional[int] = None, **opts: Any
 23    ) -> bytes:
 24        if "default" not in opts:
 25            opts["default"] = encode_numpy
 26        if ext_type_code is not None:
 27            obj_bytes = msgpack.packb(obj, use_bin_type=use_bin_type, **opts)
 28            obj_or_ext = msgpack.ExtType(ext_type_code, obj_bytes)
 29        else:
 30            obj_or_ext = obj
 31        return msgpack.packb(obj_or_ext, use_bin_type=use_bin_type, **opts)
 32
 33
 34class MsgPackDeserializer(Deserializer[bytes]):
 35    @classmethod
 36    def deserialize(
 37        cls, data: bytes, raw: bool = False, use_list: bool = False, **opts: Any
 38    ) -> Any:
 39        return msgpack.unpackb(data, raw=raw, use_list=use_list, **opts)
 40
 41
 42def to_msgpack(
 43    obj: Any,
 44    cls: Optional[Any] = None,
 45    se: type[Serializer[bytes]] = MsgPackSerializer,
 46    named: bool = True,
 47    ext_dict: Optional[dict[type[Any], int]] = None,
 48    reuse_instances: bool = False,
 49    convert_sets: bool = True,
 50    **opts: Any,
 51) -> bytes:
 52    """
 53    Serialize the object into MsgPack.
 54
 55    You can pass any serializable `obj`. If `ext_dict` option is specified, `obj` is encoded
 56    as a `msgpack.ExtType` If you supply other keyword arguments, they will be passed in
 57    `msgpack.packb` function.
 58
 59    * `named`: If `named` is True, field names are preserved, namely the object is encoded as `dict`
 60    then serialized into MsgPack.  If `named` is False, the object is encoded as `tuple` then
 61    serialized into MsgPack. `named=False` will produces compact binary.
 62
 63    * `skip_none`: When set to True, any field in the class with a None value is excluded from the
 64    serialized output. Defaults to False.
 65
 66    If you want to use the other msgpack package, you can subclass `MsgPackSerializer` and
 67    implement your own logic.
 68    """
 69    ext_type_code = None
 70    if ext_dict is not None:
 71        obj_type = type(obj)
 72        ext_type_code = ext_dict.get(obj_type)
 73        if ext_type_code is None:
 74            raise SerdeError(f"Could not find type code for {obj_type.__name__} in ext_dict")
 75
 76    kwargs: Any = {"c": cls, "reuse_instances": reuse_instances, "convert_sets": convert_sets}
 77    dict_or_tuple = to_dict(obj, **kwargs) if named else to_tuple(obj, **kwargs)
 78    return se.serialize(
 79        dict_or_tuple,
 80        ext_type_code=ext_type_code,
 81        **opts,
 82    )
 83
 84
 85@overload
 86def from_msgpack(
 87    c: type[T],
 88    s: bytes,
 89    de: type[Deserializer[bytes]] = MsgPackDeserializer,
 90    named: bool = True,
 91    ext_dict: Optional[dict[int, type[Any]]] = None,
 92    **opts: Any,
 93) -> T: ...
 94
 95
 96@overload
 97def from_msgpack(
 98    c: Any,
 99    s: bytes,
100    de: type[Deserializer[bytes]] = MsgPackDeserializer,
101    named: bool = True,
102    ext_dict: Optional[dict[int, type[Any]]] = None,
103    **opts: Any,
104) -> Any: ...
105
106
107def from_msgpack(
108    c: Any,
109    s: bytes,
110    de: type[Deserializer[bytes]] = MsgPackDeserializer,
111    named: bool = True,
112    ext_dict: Optional[dict[int, type[Any]]] = None,
113    skip_none: bool = False,
114    **opts: Any,
115) -> Any:
116    """
117    Deserialize from MsgPack into the object.
118
119    `c` is a class obejct and `s` is MsgPack binary. If `ext_dict` option is specified,
120    `c` is ignored and type is inferred from `msgpack.ExtType` If you supply other keyword
121    arguments, they will be passed in `msgpack.unpackb` function.
122
123    If you want to use the other msgpack package, you can subclass `MsgPackDeserializer`
124    and implement your own logic.
125    """
126    if ext_dict is not None:
127        ext = de.deserialize(s, **opts)
128        ext_type = ext_dict.get(ext.code)
129        if ext_type is None:
130            raise SerdeError(f"Could not find type for code {ext.code} in ext_dict")
131        return from_msgpack(ext_type, ext.data, de, named, **opts)
132    else:
133        from_func = from_dict if named else from_tuple
134        return from_func(c, de.deserialize(s, **opts), reuse_instances=False)