MessagePack Utilities

msgpack.py

import msgpack
import gc
import datetime
from datetime import timezone, timedelta

empty_dict_bytes = msgpack.packb({})

_datetime_ExtType = 42

# NOTE: if we were to store the timestamp instead of the extended information, for naive datetimes we'd have to convert
#       the object to UTC which would rely on us making an assumption about the datetime's timezone


# uses Extended types to handle datetimes: https://github.com/msgpack/msgpack-python#extended-types
def _unpacker_hook(code, data):
    if code == _datetime_ExtType:
        values = unpack(data)

        if len(values) == 8:  # we have timezone
            return datetime.datetime(*values[:-1], timezone(timedelta(seconds=values[-1])))
        else:
            return datetime.datetime(*values)

    return msgpack.ExtType(code, data)


# This will only get called for unknown types
def _packer_unknown_handler(obj):
    if isinstance(obj, datetime.datetime):
        if obj.tzinfo:
            components = (obj.year, obj.month, obj.day, obj.hour, obj.minute, obj.second, obj.microsecond, int(obj.utcoffset().total_seconds()))
        else:
            components = (obj.year, obj.month, obj.day, obj.hour, obj.minute, obj.second, obj.microsecond)

        # we double pack the values to "compress" them
        data = msgpack.ExtType(_datetime_ExtType, pack(components))
        return data

    raise TypeError("Unknown type: {}".format(type(obj)))


def pack(obj, **kwargs):
    # we don't use a global packer because it wouldn't be re-entrant safe
    return msgpack.packb(obj, use_bin_type=True, default=_packer_unknown_handler, **kwargs)


def unpack(payload: bytes, use_list=False):
    try:
        # we temporarily disable gc during unpack to bump up perf: https://pypi.python.org/pypi/msgpack-python
        gc.disable()
        # This must match the above _packer parameters above.  NOTE: use_list=False is faster
        return msgpack.unpackb(payload, use_list=use_list, raw=False, ext_hook=_unpacker_hook)
    finally:
        gc.enable()