Skip to content

Commit

Permalink
add support for ext handlers to pack
Browse files Browse the repository at this point in the history
  • Loading branch information
@vsergeev
vsergeev committed Oct 19, 2016
1 parent 6f66430 commit eb2e053
Showing 1 changed file with 83 additions and 42 deletions.
125 changes: 83 additions & 42 deletions umsgpack.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -216,7 +216,7 @@ class DuplicateKeyException(UnpackException):
# chr(obj) has a str return type instead of bytes in Python 3, and
# struct.pack(...) has the right return type in both versions.

def _pack_integer(obj, fp):
def _pack_integer(obj, fp, options):
if obj < 0:
if obj >= -32:
fp.write(struct.pack("b", obj))
Expand Down Expand Up @@ -244,19 +244,19 @@ def _pack_integer(obj, fp):
else:
raise UnsupportedTypeException("huge unsigned int")

def _pack_nil(obj, fp):
def _pack_nil(obj, fp, options):
fp.write(b"\xc0")

def _pack_boolean(obj, fp):
def _pack_boolean(obj, fp, options):
fp.write(b"\xc3" if obj else b"\xc2")

def _pack_float(obj, fp):
def _pack_float(obj, fp, options):
if _float_size == 64:
fp.write(b"\xcb" + struct.pack(">d", obj))
else:
fp.write(b"\xca" + struct.pack(">f", obj))

def _pack_string(obj, fp):
def _pack_string(obj, fp, options):
obj = obj.encode('utf-8')
if len(obj) <= 31:
fp.write(struct.pack("B", 0xa0 | len(obj)) + obj)
Expand All @@ -269,7 +269,7 @@ def _pack_string(obj, fp):
else:
raise UnsupportedTypeException("huge string")

def _pack_binary(obj, fp):
def _pack_binary(obj, fp, options):
if len(obj) <= 2**8-1:
fp.write(b"\xc4" + struct.pack("B", len(obj)) + obj)
elif len(obj) <= 2**16-1:
Expand All @@ -279,7 +279,7 @@ def _pack_binary(obj, fp):
else:
raise UnsupportedTypeException("huge binary string")

def _pack_oldspec_raw(obj, fp):
def _pack_oldspec_raw(obj, fp, options):
if len(obj) <= 31:
fp.write(struct.pack("B", 0xa0 | len(obj)) + obj)
elif len(obj) <= 2**16-1:
Expand All @@ -289,7 +289,7 @@ def _pack_oldspec_raw(obj, fp):
else:
raise UnsupportedTypeException("huge raw string")

def _pack_ext(obj, fp):
def _pack_ext(obj, fp, options):
if len(obj.data) == 1:
fp.write(b"\xd4" + struct.pack("B", obj.type & 0xff) + obj.data)
elif len(obj.data) == 2:
Expand All @@ -309,7 +309,7 @@ def _pack_ext(obj, fp):
else:
raise UnsupportedTypeException("huge ext data")

def _pack_array(obj, fp):
def _pack_array(obj, fp, options):
if len(obj) <= 15:
fp.write(struct.pack("B", 0x90 | len(obj)))
elif len(obj) <= 2**16-1:
Expand All @@ -320,9 +320,9 @@ def _pack_array(obj, fp):
raise UnsupportedTypeException("huge array")

for e in obj:
pack(e, fp)
pack(e, fp, **options)

def _pack_map(obj, fp):
def _pack_map(obj, fp, options):
if len(obj) <= 15:
fp.write(struct.pack("B", 0x80 | len(obj)))
elif len(obj) <= 2**16-1:
Expand All @@ -333,20 +333,25 @@ def _pack_map(obj, fp):
raise UnsupportedTypeException("huge array")

for k,v in obj.items():
pack(k, fp)
pack(v, fp)
pack(k, fp, **options)
pack(v, fp, **options)

########################################

# Pack for Python 2, with 'unicode' type, 'str' type, and 'long' type
def _pack2(obj, fp):
def _pack2(obj, fp, **options):
"""
Serialize a Python object into MessagePack bytes.
Args:
obj: a Python object
fp: a .write()-supporting file-like object
Kwargs:
ext_handlers (dict): dictionary of Ext handlers, mapping a custom type
to a callable that packs an instance of the type
into an Ext object
Returns:
None.
Expand All @@ -359,43 +364,58 @@ def _pack2(obj, fp):
>>> umsgpack.pack({u"compact": True, u"schema": 0}, f)
>>>
"""

global compatibility

ext_handlers = options.get("ext_handlers")

if obj is None:
_pack_nil(obj, fp)
_pack_nil(obj, fp, options)
elif ext_handlers and obj.__class__ in ext_handlers:
_pack_ext(ext_handlers[obj.__class__](obj), fp, options)
elif isinstance(obj, bool):
_pack_boolean(obj, fp)
_pack_boolean(obj, fp, options)
elif isinstance(obj, int) or isinstance(obj, long):
_pack_integer(obj, fp)
_pack_integer(obj, fp, options)
elif isinstance(obj, float):
_pack_float(obj, fp)
_pack_float(obj, fp, options)
elif compatibility and isinstance(obj, unicode):
_pack_oldspec_raw(bytes(obj), fp)
_pack_oldspec_raw(bytes(obj), fp, options)
elif compatibility and isinstance(obj, bytes):
_pack_oldspec_raw(obj, fp)
_pack_oldspec_raw(obj, fp, options)
elif isinstance(obj, unicode):
_pack_string(obj, fp)
_pack_string(obj, fp, options)
elif isinstance(obj, str):
_pack_binary(obj, fp)
_pack_binary(obj, fp, options)
elif isinstance(obj, list) or isinstance(obj, tuple):
_pack_array(obj, fp)
_pack_array(obj, fp, options)
elif isinstance(obj, dict):
_pack_map(obj, fp)
_pack_map(obj, fp, options)
elif isinstance(obj, Ext):
_pack_ext(obj, fp)
_pack_ext(obj, fp, options)
elif ext_handlers:
# Linear search for superclass
t = next((t for t in ext_handlers.keys() if isinstance(obj, t)), None)
if t:
_pack_ext(ext_handlers[t](obj), fp, options)
else:
raise UnsupportedTypeException("unsupported type: %s" % str(type(obj)))
else:
raise UnsupportedTypeException("unsupported type: %s" % str(type(obj)))

# Pack for Python 3, with unicode 'str' type, 'bytes' type, and no 'long' type
def _pack3(obj, fp):
def _pack3(obj, fp, **options):
"""
Serialize a Python object into MessagePack bytes.
Args:
obj: a Python object
fp: a .write()-supporting file-like object
Kwargs:
ext_handlers (dict): dictionary of Ext handlers, mapping a custom type
to a callable that packs an instance of the type
into an Ext object
Returns:
None.
Expand All @@ -410,38 +430,54 @@ def _pack3(obj, fp):
"""
global compatibility

ext_handlers = options.get("ext_handlers")

if obj is None:
_pack_nil(obj, fp)
_pack_nil(obj, fp, options)
elif ext_handlers and obj.__class__ in ext_handlers:
_pack_ext(ext_handlers[obj.__class__](obj), fp, options)
elif isinstance(obj, bool):
_pack_boolean(obj, fp)
_pack_boolean(obj, fp, options)
elif isinstance(obj, int):
_pack_integer(obj, fp)
_pack_integer(obj, fp, options)
elif isinstance(obj, float):
_pack_float(obj, fp)
_pack_float(obj, fp, options)
elif compatibility and isinstance(obj, str):
_pack_oldspec_raw(obj.encode('utf-8'), fp)
_pack_oldspec_raw(obj.encode('utf-8'), fp, options)
elif compatibility and isinstance(obj, bytes):
_pack_oldspec_raw(obj, fp)
_pack_oldspec_raw(obj, fp, options)
elif isinstance(obj, str):
_pack_string(obj, fp)
_pack_string(obj, fp, options)
elif isinstance(obj, bytes):
_pack_binary(obj, fp)
_pack_binary(obj, fp, options)
elif isinstance(obj, list) or isinstance(obj, tuple):
_pack_array(obj, fp)
_pack_array(obj, fp, options)
elif isinstance(obj, dict):
_pack_map(obj, fp)
_pack_map(obj, fp, options)
elif isinstance(obj, Ext):
_pack_ext(obj, fp)
_pack_ext(obj, fp, options)
elif ext_handlers:
# Linear search for superclass
t = next((t for t in ext_handlers.keys() if isinstance(obj, t)), None)
if t:
_pack_ext(ext_handlers[t](obj), fp, options)
else:
raise UnsupportedTypeException("unsupported type: %s" % str(type(obj)))
else:
raise UnsupportedTypeException("unsupported type: %s" % str(type(obj)))

def _packb2(obj):
def _packb2(obj, **options):
"""
Serialize a Python object into MessagePack bytes.
Args:
obj: a Python object
Kwargs:
ext_handlers (dict): dictionary of Ext handlers, mapping a custom type
to a callable that packs an instance of the type
into an Ext object
Returns:
A 'str' containing serialized MessagePack bytes.
Expand All @@ -455,16 +491,21 @@ def _packb2(obj):
>>>
"""
fp = io.BytesIO()
_pack2(obj, fp)
_pack2(obj, fp, **options)
return fp.getvalue()

def _packb3(obj):
def _packb3(obj, **options):
"""
Serialize a Python object into MessagePack bytes.
Args:
obj: a Python object
Kwargs:
ext_handlers (dict): dictionary of Ext handlers, mapping a custom type
to a callable that packs an instance of the type
into an Ext object
Returns:
A 'bytes' containing serialized MessagePack bytes.
Expand All @@ -478,7 +519,7 @@ def _packb3(obj):
>>>
"""
fp = io.BytesIO()
_pack3(obj, fp)
_pack3(obj, fp, **options)
return fp.getvalue()

################################################################################
Expand Down

0 comments on commit eb2e053

Please sign in to comment.