WAMP serialization module

serialization.py

import codecs
import json
import sys
from base64 import b64decode, b64encode
from codecs import CodecInfo
from functools import partial
from json import JSONDecoder, scanner

import six

from wampproto.exceptions import SerializationError

try:
    import msgpack
except ImportError:
    try:
        import umsgpack as msgpack
    except ImportError:
        msgpack = None

try:
    import cbor
except ImportError:
    try:
        import cbor2 as cbor
    except ImportError:
        cbor = None

try:
    import ubjson
except ImportError:
    ubjson = None


def serialize(obj, serializer):
    try:
        serialize = serializers[serializer]
    except KeyError:
        six.raise_from(SerializationError('no such serializer: {}'.format(serializer), obj), None)

    try:
        return serialize(obj)
    except Exception as e:
        six.raise_from(SerializationError('error serializing message', obj), e)


def deserialize(data, serializer):
    try:
        deserialize = deserializers[serializer]
    except KeyError:
        six.raise_from(SerializationError('no such serializer: {}'.format(serializer), data), None)

    try:
        return deserialize(data)
    except Exception as e:
        six.raise_from(SerializationError('error deserializing message', data), e)


class BinaryStringDecoder(JSONDecoder):
    def __init__(self, *args, **kwargs):
        super(BinaryStringDecoder, self).__init__(*args, **kwargs)
        self.parse_normal_string = self.parse_string
        self.parse_string = self.parse_binary_string
        self.scan_once = scanner.py_make_scanner(self)

    def parse_binary_string(self, *args, **kwargs):
        string, end = self.parse_normal_string(*args, **kwargs)
        if string.startswith(u'\x00'):
            return b64decode(string[1:].encode('ascii')), end
        else:
            return string, end


def decode_bytes_wamp(input, errors='strict'):
    return u'\x00' + b64encode(input), len(input)


def json_binary_default(obj):
    if isinstance(obj, (bytes, bytearray)):
        return u'\x00' + b64encode(obj).decode('ascii')

    raise TypeError('cannot serialize type {} to JSON'.format(obj.__class__.__name__))


# Apply special WAMP hacks to (de)serialize binary data with JSON
if sys.version_info[0] < 3:
    json_encoding = 'wamp'
    utf8_info = codecs.lookup('utf-8')
    wamp_utf8_info = CodecInfo(utf8_info.encode, decode_bytes_wamp, name='wamp')
    codecs.register(lambda encoding: wamp_utf8_info if encoding == json_encoding else None)
    json_dumps_kwargs = {'encoding': json_encoding}
else:
    json_dumps_kwargs = {'default': json_binary_default}


serializers = {'json': partial(json.dumps, **json_dumps_kwargs)}
deserializers = {'json': partial(json.loads, cls=BinaryStringDecoder)}

if msgpack:
    serializers['msgpack'] = msgpack.packb
    deserializers['msgpack'] = msgpack.unpackb
if cbor:
    serializers['cbor'] = cbor.dumps
    deserializers['cbor'] = cbor.loads
if ubjson:
    serializers['ubjson'] = msgpack.packb
    deserializers['ubjson'] = msgpack.unpackb