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