8000 Pure Python packer supports memoryview of multi byte items. · wumingcp-ps/msgpack-python@fc29338 · GitHub
[go: up one dir, main page]

Skip to content

Commit fc29338

Browse files
committed
Pure Python packer supports memoryview of multi byte items.
2 parents ceb9635 + 53f47ef commit fc29338

File tree

3 files changed

+159
-41
lines changed

3 files changed

+159
-41
lines changed

msgpack/_packer.pyx

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ cdef extern from "pack.h":
3939
int msgpack_pack_ext(msgpack_packer* pk, char typecode, size_t l)
4040

4141
cdef int DEFAULT_RECURSE_LIMIT=511
42+
cdef size_t ITEM_LIMIT = (2**32)-1
4243

4344

4445
cdef class Packer(object):
@@ -178,7 +179,7 @@ cdef class Packer(object):
178179
ret = msgpack_pack_double(&self.pk, dval)
179180
elif PyBytes_CheckExact(o) if strict_types else PyBytes_Check(o):
180181
L = len(o)
181-
if L > (2**32)-1:
182+
if L > ITEM_LIMIT:
182183
raise PackValueError("bytes is too large")
183184
rawval = o
184185
ret = msgpack_pack_bin(&self.pk, L)
@@ -189,7 +190,7 @@ cdef class Packer(object):
189190
raise TypeError("Can't encode unicode string: no encoding is specified")
190191
o = PyUnicode_AsEncodedString(o, self.encoding, self.unicode_errors)
191192
L = len(o)
192-
if L > (2**32)-1:
193+
if L > ITEM_LIMIT:
193194
raise PackValueError("unicode string is too large")
194195
rawval = o
195196
ret = msgpack_pack_raw(&self.pk, L)
@@ -198,7 +199,7 @@ cdef class Packer(object):
198199
elif PyDict_CheckExact(o):
199200
d = <dict>o
200201
L = len(d)
201-
if L > (2**32)-1:
202+
if L > ITEM_LIMIT< 8000 /span>:
202203
raise PackValueError("dict is too large")
203204
ret = msgpack_pack_map(&self.pk, L)
204205
if ret == 0:
@@ -209,7 +210,7 @@ cdef class Packer(object):
209210
if ret != 0: break
210211
elif not strict_types and PyDict_Check(o):
211212
L = len(o)
212-
if L > (2**32)-1:
213+
if L > ITEM_LIMIT:
213214
raise PackValueError("dict is too large")
214215
ret = msgpack_pack_map(&self.pk, L)
215216
if ret == 0:
@@ -223,13 +224,13 @@ cdef class Packer(object):
223224
longval = o.code
224225
rawval = o.data
225226
L = len(o.data)
226-
if L > (2**32)-1:
227+
if L > ITEM_LIMIT:
227228
raise PackValueError("EXT data is too large")
228229
ret = msgpack_pack_ext(&self.pk, longval, L)
229230
ret = msgpack_pack_raw_body(&self.pk, rawval, L)
230231
elif PyList_CheckExact(o) if strict_types else (PyTuple_Check(o) or PyList_Check(o)):
231232
L = len(o)
232-
if L > (2**32)-1:
233+
if L > ITEM_LIMIT:
233234
raise PackValueError("list is too large")
234235
ret = msgpack_pack_array(&self.pk, L)
235236
if ret == 0:
@@ -240,7 +241,7 @@ cdef class Packer(object):
240241
if PyObject_GetBuffer(o, &view, PyBUF_SIMPLE) != 0:
241242
raise PackValueError("could not get buffer for memoryview")
242243
L = view.len
243-
if L > (2**32)-1:
244+
if L > ITEM_LIMIT:
244245
PyBuffer_Release(&view);
245246
raise PackValueError("memoryview is too large")
246247
ret = msgpack_pack_bin(&self.pk, L)
@@ -271,8 +272,8 @@ cdef class Packer(object):
271272
msgpack_pack_ext(&self.pk, typecode, len(data))
272273
msgpack_pack_raw_body(&self.pk, data, len(data))
273274

274-
def pack_array_header(self, long long size):
275-
if size > (2**32-1):
275+
def pack_array_header(self, size_t size):
276+
if size > ITEM_LIMIT:
276277
raise PackValueError
277278
cdef int ret = msgpack_pack_array(&self.pk, size)
278279
if ret == -1:
@@ -284,8 +285,8 @@ cdef class Packer(object):
284285
self.pk.length = 0
285286
return buf
286287

287-
def pack_map_header(self, long long size):
288-
if size > (2**32-1):
288+
def pack_map_header(self, size_t size):
289+
if size > ITEM_LIMIT:
289290
raise PackValueError
290291
cdef int ret = msgpack_pack_map(&self.pk, size)
291292
if ret == -1:

msgpack/fallback.py

Lines changed: 41 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -685,35 +685,28 @@ def _pack(self, obj, nest_limit=DEFAULT_RECURSE_LIMIT,
685685
default_used = True
686686
continue
687687
raise PackOverflowError("Integer value out of range")
688-
if self._use_bin_type and check(obj, (bytes, memoryview)):
688+
if check(obj, bytes):
689689
n = len(obj)
690-
if n <= 0xff:
691-
self._buffer.write(struct.pack('>BB', 0xc4, n))
692-
elif n <= 0xffff:
693-
self._buffer.write(struct.pack(">BH", 0xc5, n))
694-
elif n <= 0xffffffff:
695-
self._buffer.write(struct.pack(">BI", 0xc6, n))
696-
else:
690+
if n >= 2**32:
697691
raise PackValueError("Bytes is too large")
692+
self._fb_pack_bin_header(n)
698693
return self._buffer.write(obj)
699-
if check(obj, (Unicode, bytes, memoryview)):
700-
if check(obj, Unicode):
701-
if self._encoding is None:
702-
raise TypeError(
703-
"Can't encode unicode string: "
704-
"no encoding is specified")
705-
obj = obj.encode(self._encoding, self._unicode_errors)
694+
if check(obj, Unicode):
695+
if self._encoding is None:
696+
raise TypeError(
697+
"Can't encode unicode string: "
698+
"no encoding is specified")
699+
obj = obj.encode(self._encoding, self._unicode_errors)
706700
n = len(obj)
707-
if n <= 0x1f:
708-
self._buffer.write(struct.pack('B', 0xa0 + n))
709-
elif self._use_bin_type and n <= 0xff:
710-
self._buffer.write(struct.pack('>BB', 0xd9, n))
711-
elif n <= 0xffff:
712-
self._buffer.write(struct.pack(">BH", 0xda, n))
713-
elif n <= 0xffffffff:
714-
self._buffer.write(struct.pack(">BI", 0xdb, n))
715-
else:
701+
if n >= 2**32:
716702
raise PackValueError("String is too large")
703+
self._fb_pack_raw_header(n)
704+
return self._buffer.write(obj)
705+
if check(obj, memoryview):
706+
n = len(obj) * obj.itemsize
707+
if n >= 2**32:
708+
raise PackValueError("Memoryview is too large")
709+
self._fb_pack_bin_header(n)
717710
return self._buffer.write(obj)
718711
if check(obj, float):
719712
if self._use_float:
@@ -852,6 +845,30 @@ def _fb_pack_map_pairs(self, n, pairs, nest_limit=DEFAULT_RECURSE_LIMIT):
852845
self._pack(k, nest_limit - 1)
853846
self._pack(v, nest_limit - 1)
854847

848+
def _fb_pack_raw_header(self, n):
849+
if n <= 0x1f:
850+
self._buffer.write(struct.pack('B', 0xa0 + n))
851+
elif self._use_bin_type and n <= 0xff:
852+
self._buffer.write(struct.pack('>BB', 0xd9, n))
853+
elif n <= 0xffff:
854+
self._buffer.write(struct.pack(">BH", 0xda, n))
855+
elif n <= 0xffffffff:
856+
self._buffer.write(struct.pack(">BI", 0xdb, n))
857+
else:
858+
raise PackValueError('Raw is too large')
859+
860+
def _fb_pack_bin_header(self, n):
861+
if not self._use_bin_type:
862+
return self._fb_pack_raw_header(n)
863+
elif n <= 0xff:
864+
return self._buffer.write(struct.pack('>BB', 0xc4, n))
865+
elif n <= 0xffff:
866+
return self._buffer.write(struct.pack(">BH", 0xc5, n))
867+
elif n <= 0xffffffff:
868+
return self._buffer.write(struct.pack(">BI", 0xc6, n))
869+
else:
870+
raise PackValueError('Bin is too large')
871+
855872
def bytes(self):
856873
return self._buffer.getvalue()
857874

test/test_memoryview.py

Lines changed: 106 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,112 @@
11
#!/usr/bin/env python
22
# coding: utf-8
33

4-
4+
from array import array
55
from msgpack import packb, unpackb
6+
import sys
7+
8+
9+
# For Python < 3:
10+
# - array type only supports old buffer interface
11+
# - array.frombytes is not available, must use deprecated array.fromstring
12+
if sys.version_info[0] < 3:
13+
def make_memoryview(obj):
14+
return memoryview(buffer(obj))
15+
16+
def make_array(f, data):
17+
a = array(f)
18+
a.fromstring(data)
19+
return a
20+
21+
def get_data(a):
22+
return a.tostring()
23+
else:
24+
make_memoryview = memoryview
25+
26+
def make_array(f, data):
27+
a = array(f)
28+
a.frombytes(data)
29+
return a
30+
31+
def get_data(a):
32+
return a.tobytes()
33+
34+
35+
def _runtest(format, nbytes, expected_header, expected_prefix, use_bin_type):
36+
# create a new array
37+
original_array = array(format)
38+
original_array.fromlist([255] * (nbytes // original_array.itemsize))
39+
original_data = get_data(original_array)
40+
view = make_memoryview(original_array)
41+
42+
# pack, unpack, and reconstruct array
43+
packed = packb(view, use_bin_type=use_bin_type)
44+
unpacked = unpackb(packed)
45+
reconstructed_array = make_array(format, unpacked)
46+
47+
# check that we got the right amount of data
48+
assert len(original_data) == nbytes
49+
# check packed header
50+
assert packed[:1] == expected_header
51+
# check packed length prefix, if any
52+
assert packed[1:1+len(expected_prefix)] == expected_prefix
53+
# check packed data
54+
assert packed[1+len(expected_prefix):] == original_data
55+
# check array unpacked correctly
56+
assert original_array == reconstructed_array
57+
58+
59+
def test_fixstr_from_byte():
60+
_runtest('B', 1, b'\xa1', b'', False)
61+
_runtest('B', 31, b'\xbf', b'', False)
62+
63+
64+
def test_fixstr_from_float():
65+
_runtest('f', 4, b'\xa4', b'', False)
66+
_runtest('f', 28, b'\xbc', b'', False)
67+
68+
69+
def test_str16_from_byte():
70+
_runtest('B', 2**8, b'\xda', b'\x01\x00', False)
71+
_runtest('B', 2**16-1, b'\xda', b'\xff\xff', False)
72+
73+
74+
def test_str16_from_float():
75+
_runtest('f', 2**8, b'\xda', b'\x01\x00', False)
76+
_runtest('f', 2**16-4, b'\xda', b'\xff\xfc', False)
77+
78+
79+
def test_str32_from_byte():
80+
_runtest('B', 2**16, b'\xdb', b'\x00\x01\x00\x00', False)
81+
82+
83+
def test_str32_from_float():
84+
_runtest('f', 2**16, b'\xdb', b'\x00\x01\x00\x00', False)
85+
86+
87+
def test_bin8_from_byte():
88+
_runtest('B', 1, b'\xc4', b'\x01', True)
89+
_runtest('B', 2**8-1, b'\xc4', b'\xff', True)
90+
91+
92+
def test_bin8_from_float():
93+
_runtest('f', 4, b'\xc4', b'\x04', True)
94+
_runtest('f', 2**8-4, b'\xc4', b'\xfc', True)
95+
96+
97+
def test_bin16_from_byte():
98+
_runtest('B', 2**8, b'\xc5', b'\x01\x00', True)
99+
_runtest('B', 2**16-1, b'\xc5', b'\xff\xff', True)
100+
101+
102+
def test_bin16_from_float():
103+
_runtest('f', 2**8, b'\xc5', b'\x01\x00', True)
104+
_runtest('f', 2**16-4, b'\xc5', b'\xff\xfc', True)
105+
106+
107+
def test_bin32_from_byte():
108+
_runtest('B', 2**16, b'\xc6', b'\x00\x01\x00\x00', True)
6109

7110

8-
def test_pack_memoryview():
9-
data = bytearray(range(256))
10-
view = memoryview(data)
11-
unpacked = unpackb(packb(view))
12-
assert data == unpacked
111+
def test_bin32_from_float():
112+
_runtest('f', 2**16, b'\xc6', b'\x00\x01\x00\x00', True)

0 commit comments

Comments
 (0)
0