Pure Python packer supports memoryview of multi byte items.
This commit is contained in:
INADA Naoki
commit
fc2933853a
|
|
@ -39,6 +39,7 @@ cdef extern from "pack.h":
|
|||
int msgpack_pack_ext(msgpack_packer* pk, char typecode, size_t l)
|
||||
|
||||
cdef int DEFAULT_RECURSE_LIMIT=511
|
||||
cdef size_t ITEM_LIMIT = (2**32)-1
|
||||
|
||||
|
||||
cdef class Packer(object):
|
||||
|
|
@ -178,7 +179,7 @@ cdef class Packer(object):
|
|||
ret = msgpack_pack_double(&self.pk, dval)
|
||||
elif PyBytes_CheckExact(o) if strict_types else PyBytes_Check(o):
|
||||
L = len(o)
|
||||
if L > (2**32)-1:
|
||||
if L > ITEM_LIMIT:
|
||||
raise PackValueError("bytes is too large")
|
||||
rawval = o
|
||||
ret = msgpack_pack_bin(&self.pk, L)
|
||||
|
|
@ -189,7 +190,7 @@ cdef class Packer(object):
|
|||
raise TypeError("Can't encode unicode string: no encoding is specified")
|
||||
o = PyUnicode_AsEncodedString(o, self.encoding, self.unicode_errors)
|
||||
L = len(o)
|
||||
if L > (2**32)-1:
|
||||
if L > ITEM_LIMIT:
|
||||
raise PackValueError("unicode string is too large")
|
||||
rawval = o
|
||||
ret = msgpack_pack_raw(&self.pk, L)
|
||||
|
|
@ -198,7 +199,7 @@ cdef class Packer(object):
|
|||
elif PyDict_CheckExact(o):
|
||||
d = <dict>o
|
||||
L = len(d)
|
||||
if L > (2**32)-1:
|
||||
if L > ITEM_LIMIT:
|
||||
raise PackValueError("dict is too large")
|
||||
ret = msgpack_pack_map(&self.pk, L)
|
||||
if ret == 0:
|
||||
|
|
@ -209,7 +210,7 @@ cdef class Packer(object):
|
|||
if ret != 0: break
|
||||
elif not strict_types and PyDict_Check(o):
|
||||
L = len(o)
|
||||
if L > (2**32)-1:
|
||||
if L > ITEM_LIMIT:
|
||||
raise PackValueError("dict is too large")
|
||||
ret = msgpack_pack_map(&self.pk, L)
|
||||
if ret == 0:
|
||||
|
|
@ -223,13 +224,13 @@ cdef class Packer(object):
|
|||
longval = o.code
|
||||
rawval = o.data
|
||||
L = len(o.data)
|
||||
if L > (2**32)-1:
|
||||
if L > ITEM_LIMIT:
|
||||
raise PackValueError("EXT data is too large")
|
||||
ret = msgpack_pack_ext(&self.pk, longval, L)
|
||||
ret = msgpack_pack_raw_body(&self.pk, rawval, L)
|
||||
elif PyList_CheckExact(o) if strict_types else (PyTuple_Check(o) or PyList_Check(o)):
|
||||
L = len(o)
|
||||
if L > (2**32)-1:
|
||||
if L > ITEM_LIMIT:
|
||||
raise PackValueError("list is too large")
|
||||
ret = msgpack_pack_array(&self.pk, L)
|
||||
if ret == 0:
|
||||
|
|
@ -240,7 +241,7 @@ cdef class Packer(object):
|
|||
if PyObject_GetBuffer(o, &view, PyBUF_SIMPLE) != 0:
|
||||
raise PackValueError("could not get buffer for memoryview")
|
||||
L = view.len
|
||||
if L > (2**32)-1:
|
||||
if L > ITEM_LIMIT:
|
||||
PyBuffer_Release(&view);
|
||||
raise PackValueError("memoryview is too large")
|
||||
ret = msgpack_pack_bin(&self.pk, L)
|
||||
|
|
@ -271,8 +272,8 @@ cdef class Packer(object):
|
|||
msgpack_pack_ext(&self.pk, typecode, len(data))
|
||||
msgpack_pack_raw_body(&self.pk, data, len(data))
|
||||
|
||||
def pack_array_header(self, long long size):
|
||||
if size > (2**32-1):
|
||||
def pack_array_header(self, size_t size):
|
||||
if size > ITEM_LIMIT:
|
||||
raise PackValueError
|
||||
cdef int ret = msgpack_pack_array(&self.pk, size)
|
||||
if ret == -1:
|
||||
|
|
@ -284,8 +285,8 @@ cdef class Packer(object):
|
|||
self.pk.length = 0
|
||||
return buf
|
||||
|
||||
def pack_map_header(self, long long size):
|
||||
if size > (2**32-1):
|
||||
def pack_map_header(self, size_t size):
|
||||
if size > ITEM_LIMIT:
|
||||
raise PackValueError
|
||||
cdef int ret = msgpack_pack_map(&self.pk, size)
|
||||
if ret == -1:
|
||||
|
|
|
|||
|
|
@ -685,35 +685,28 @@ class Packer(object):
|
|||
default_used = True
|
||||
continue
|
||||
raise PackOverflowError("Integer value out of range")
|
||||
if self._use_bin_type and check(obj, (bytes, memoryview)):
|
||||
if check(obj, bytes):
|
||||
n = len(obj)
|
||||
if n <= 0xff:
|
||||
self._buffer.write(struct.pack('>BB', 0xc4, n))
|
||||
elif n <= 0xffff:
|
||||
self._buffer.write(struct.pack(">BH", 0xc5, n))
|
||||
elif n <= 0xffffffff:
|
||||
self._buffer.write(struct.pack(">BI", 0xc6, n))
|
||||
else:
|
||||
if n >= 2**32:
|
||||
raise PackValueError("Bytes is too large")
|
||||
self._fb_pack_bin_header(n)
|
||||
return self._buffer.write(obj)
|
||||
if check(obj, (Unicode, bytes, memoryview)):
|
||||
if check(obj, Unicode):
|
||||
if self._encoding is None:
|
||||
raise TypeError(
|
||||
"Can't encode unicode string: "
|
||||
"no encoding is specified")
|
||||
obj = obj.encode(self._encoding, self._unicode_errors)
|
||||
if check(obj, Unicode):
|
||||
if self._encoding is None:
|
||||
raise TypeError(
|
||||
"Can't encode unicode string: "
|
||||
"no encoding is specified")
|
||||
obj = obj.encode(self._encoding, self._unicode_errors)
|
||||
n = len(obj)
|
||||
if n <= 0x1f:
|
||||
self._buffer.write(struct.pack('B', 0xa0 + n))
|
||||
elif self._use_bin_type and n <= 0xff:
|
||||
self._buffer.write(struct.pack('>BB', 0xd9, n))
|
||||
elif n <= 0xffff:
|
||||
self._buffer.write(struct.pack(">BH", 0xda, n))
|
||||
elif n <= 0xffffffff:
|
||||
self._buffer.write(struct.pack(">BI", 0xdb, n))
|
||||
else:
|
||||
if n >= 2**32:
|
||||
raise PackValueError("String is too large")
|
||||
self._fb_pack_raw_header(n)
|
||||
return self._buffer.write(obj)
|
||||
if check(obj, memoryview):
|
||||
n = len(obj) * obj.itemsize
|
||||
if n >= 2**32:
|
||||
raise PackValueError("Memoryview is too large")
|
||||
self._fb_pack_bin_header(n)
|
||||
return self._buffer.write(obj)
|
||||
if check(obj, float):
|
||||
if self._use_float:
|
||||
|
|
@ -852,6 +845,30 @@ class Packer(object):
|
|||
self._pack(k, nest_limit - 1)
|
||||
self._pack(v, nest_limit - 1)
|
||||
|
||||
def _fb_pack_raw_header(self, n):
|
||||
if n <= 0x1f:
|
||||
self._buffer.write(struct.pack('B', 0xa0 + n))
|
||||
elif self._use_bin_type and n <= 0xff:
|
||||
self._buffer.write(struct.pack('>BB', 0xd9, n))
|
||||
elif n <= 0xffff:
|
||||
self._buffer.write(struct.pack(">BH", 0xda, n))
|
||||
elif n <= 0xffffffff:
|
||||
self._buffer.write(struct.pack(">BI", 0xdb, n))
|
||||
else:
|
||||
raise PackValueError('Raw is too large')
|
||||
|
||||
def _fb_pack_bin_header(self, n):
|
||||
if not self._use_bin_type:
|
||||
return self._fb_pack_raw_header(n)
|
||||
elif n <= 0xff:
|
||||
return self._buffer.write(struct.pack('>BB', 0xc4, n))
|
||||
elif n <= 0xffff:
|
||||
return self._buffer.write(struct.pack(">BH", 0xc5, n))
|
||||
elif n <= 0xffffffff:
|
||||
return self._buffer.write(struct.pack(">BI", 0xc6, n))
|
||||
else:
|
||||
raise PackValueError('Bin is too large')
|
||||
|
||||
def bytes(self):
|
||||
return self._buffer.getvalue()
|
||||
|
||||
|
|
|
|||
|
|
@ -1,12 +1,112 @@
|
|||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
|
||||
|
||||
from array import array
|
||||
from msgpack import packb, unpackb
|
||||
import sys
|
||||
|
||||
|
||||
def test_pack_memoryview():
|
||||
data = bytearray(range(256))
|
||||
view = memoryview(data)
|
||||
unpacked = unpackb(packb(view))
|
||||
assert data == unpacked
|
||||
# For Python < 3:
|
||||
# - array type only supports old buffer interface
|
||||
# - array.frombytes is not available, must use deprecated array.fromstring
|
||||
if sys.version_info[0] < 3:
|
||||
def make_memoryview(obj):
|
||||
return memoryview(buffer(obj))
|
||||
|
||||
def make_array(f, data):
|
||||
a = array(f)
|
||||
a.fromstring(data)
|
||||
return a
|
||||
|
||||
def get_data(a):
|
||||
return a.tostring()
|
||||
else:
|
||||
make_memoryview = memoryview
|
||||
|
||||
def make_array(f, data):
|
||||
a = array(f)
|
||||
a.frombytes(data)
|
||||
return a
|
||||
|
||||
def get_data(a):
|
||||
return a.tobytes()
|
||||
|
||||
|
||||
def _runtest(format, nbytes, expected_header, expected_prefix, use_bin_type):
|
||||
# create a new array
|
||||
original_array = array(format)
|
||||
original_array.fromlist([255] * (nbytes // original_array.itemsize))
|
||||
original_data = get_data(original_array)
|
||||
view = make_memoryview(original_array)
|
||||
|
||||
# pack, unpack, and reconstruct array
|
||||
packed = packb(view, use_bin_type=use_bin_type)
|
||||
unpacked = unpackb(packed)
|
||||
reconstructed_array = make_array(format, unpacked)
|
||||
|
||||
# check that we got the right amount of data
|
||||
assert len(original_data) == nbytes
|
||||
# check packed header
|
||||
assert packed[:1] == expected_header
|
||||
# check packed length prefix, if any
|
||||
assert packed[1:1+len(expected_prefix)] == expected_prefix
|
||||
# check packed data
|
||||
assert packed[1+len(expected_prefix):] == original_data
|
||||
# check array unpacked correctly
|
||||
assert original_array == reconstructed_array
|
||||
|
||||
|
||||
def test_fixstr_from_byte():
|
||||
_runtest('B', 1, b'\xa1', b'', False)
|
||||
_runtest('B', 31, b'\xbf', b'', False)
|
||||
|
||||
|
||||
def test_fixstr_from_float():
|
||||
_runtest('f', 4, b'\xa4', b'', False)
|
||||
_runtest('f', 28, b'\xbc', b'', False)
|
||||
|
||||
|
||||
def test_str16_from_byte():
|
||||
_runtest('B', 2**8, b'\xda', b'\x01\x00', False)
|
||||
_runtest('B', 2**16-1, b'\xda', b'\xff\xff', False)
|
||||
|
||||
|
||||
def test_str16_from_float():
|
||||
_runtest('f', 2**8, b'\xda', b'\x01\x00', False)
|
||||
_runtest('f', 2**16-4, b'\xda', b'\xff\xfc', False)
|
||||
|
||||
|
||||
def test_str32_from_byte():
|
||||
_runtest('B', 2**16, b'\xdb', b'\x00\x01\x00\x00', False)
|
||||
|
||||
|
||||
def test_str32_from_float():
|
||||
_runtest('f', 2**16, b'\xdb', b'\x00\x01\x00\x00', False)
|
||||
|
||||
|
||||
def test_bin8_from_byte():
|
||||
_runtest('B', 1, b'\xc4', b'\x01', True)
|
||||
_runtest('B', 2**8-1, b'\xc4', b'\xff', True)
|
||||
|
||||
|
||||
def test_bin8_from_float():
|
||||
_runtest('f', 4, b'\xc4', b'\x04', True)
|
||||
_runtest('f', 2**8-4, b'\xc4', b'\xfc', True)
|
||||
|
||||
|
||||
def test_bin16_from_byte():
|
||||
_runtest('B', 2**8, b'\xc5', b'\x01\x00', True)
|
||||
_runtest('B', 2**16-1, b'\xc5', b'\xff\xff', True)
|
||||
|
||||
|
||||
def test_bin16_from_float():
|
||||
_runtest('f', 2**8, b'\xc5', b'\x01\x00', True)
|
||||
_runtest('f', 2**16-4, b'\xc5', b'\xff\xfc', True)
|
||||
|
||||
|
||||
def test_bin32_from_byte():
|
||||
_runtest('B', 2**16, b'\xc6', b'\x00\x01\x00\x00', True)
|
||||
|
||||
|
||||
def test_bin32_from_float():
|
||||
_runtest('f', 2**16, b'\xc6', b'\x00\x01\x00\x00', True)
|
||||
|
|
|
|||
Loading…
Reference in New Issue