1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93
|
# coding: utf-8
import functools
import catalogue
# These need to be imported before packer and unpacker
from ._epoch import utc, epoch # noqa
from ._version import version
from .exceptions import *
# In msgpack-python these are put under a _cmsgpack module that textually includes
# them. I dislike this so I refactored it.
from ._packer import Packer as _Packer
from ._unpacker import unpackb as _unpackb
from ._unpacker import Unpacker as _Unpacker
from .ext 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
data = stream.read()
return _unpackb(data, **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
|