88 lines
2.5 KiB
Python
88 lines
2.5 KiB
Python
# coding: utf-8
|
|
|
|
import functools
|
|
import catalogue
|
|
|
|
from ._version import version
|
|
from .exceptions import *
|
|
from ._packer import Packer as _Packer
|
|
from ._unpacker import unpackb as _unpackb
|
|
from ._unpacker import unpack as _unpack
|
|
from ._unpacker import Unpacker as _Unpacker
|
|
from ._ext_type import ExtType
|
|
from ._msgpack_numpy import encode_numpy as _encode_numpy
|
|
from ._msgpack_numpy import decode_numpy as _decode_numpy
|
|
|
|
|
|
msgpack_encoders = catalogue.create("srsly", "msgpack_encoders", entry_points=True)
|
|
msgpack_decoders = catalogue.create("srsly", "msgpack_decoders", entry_points=True)
|
|
|
|
msgpack_encoders.register("numpy", func=_encode_numpy)
|
|
msgpack_decoders.register("numpy", func=_decode_numpy)
|
|
|
|
|
|
# msgpack_numpy extensions
|
|
class Packer(_Packer):
|
|
def __init__(self, *args, **kwargs):
|
|
default = kwargs.get("default")
|
|
for encoder in msgpack_encoders.get_all().values():
|
|
default = functools.partial(encoder, chain=default)
|
|
kwargs["default"] = default
|
|
super(Packer, self).__init__(*args, **kwargs)
|
|
|
|
|
|
class Unpacker(_Unpacker):
|
|
def __init__(self, *args, **kwargs):
|
|
object_hook = kwargs.get("object_hook")
|
|
for decoder in msgpack_decoders.get_all().values():
|
|
object_hook = functools.partial(decoder, chain=object_hook)
|
|
kwargs["object_hook"] = object_hook
|
|
super(Unpacker, self).__init__(*args, **kwargs)
|
|
|
|
|
|
def pack(o, stream, **kwargs):
|
|
"""
|
|
Pack an object and write it to a stream.
|
|
"""
|
|
packer = Packer(**kwargs)
|
|
stream.write(packer.pack(o))
|
|
|
|
|
|
def packb(o, **kwargs):
|
|
"""
|
|
Pack an object and return the packed bytes.
|
|
"""
|
|
return Packer(**kwargs).pack(o)
|
|
|
|
|
|
def unpack(stream, **kwargs):
|
|
"""
|
|
Unpack a packed object from a stream.
|
|
"""
|
|
if "object_pairs_hook" not in kwargs:
|
|
object_hook = kwargs.get("object_hook")
|
|
for decoder in msgpack_decoders.get_all().values():
|
|
object_hook = functools.partial(decoder, chain=object_hook)
|
|
kwargs["object_hook"] = object_hook
|
|
return _unpack(stream, **kwargs)
|
|
|
|
|
|
def unpackb(packed, **kwargs):
|
|
"""
|
|
Unpack a packed object.
|
|
"""
|
|
if "object_pairs_hook" not in kwargs:
|
|
object_hook = kwargs.get("object_hook")
|
|
for decoder in msgpack_decoders.get_all().values():
|
|
object_hook = functools.partial(decoder, chain=object_hook)
|
|
kwargs["object_hook"] = object_hook
|
|
return _unpackb(packed, **kwargs)
|
|
|
|
|
|
# alias for compatibility to simplejson/marshal/pickle.
|
|
load = unpack
|
|
loads = unpackb
|
|
|
|
dump = pack
|
|
dumps = packb
|