Skip to content

Commit fd61226

Browse files
nightsailernightsailer
committed
Merge pull request #2 from jfolz/master
merge memoryview staff from jfolz
2 parents 9b9cbaa + e47e2e5 commit fd61226

File tree

3 files changed

+74
-30
lines changed

3 files changed

+74
-30
lines changed

‎msgpack/_unpacker.pyx‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,14 @@ from cpython.bytes cimport (
88
)
99
from cpython.buffer cimport (
1010
Py_buffer,
11+
PyObject_CheckBuffer,
1112
PyBuffer_Release,
1213
PyObject_GetBuffer,
1314
PyBUF_SIMPLE,
1415
)
1516
from cpython.mem cimport PyMem_Malloc, PyMem_Free
1617
from cpython.object cimport PyCallable_Check
18+
from cpython.exc cimport PyErr_WarnEx
1719

1820
cdef extern from "Python.h":
1921
ctypedef struct PyObject
@@ -129,12 +131,27 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
129131
cdef Py_ssize_t off = 0
130132
cdef int ret
131133

132-
cdef char* buf
134+
cdef Py_buffer view
135+
cdef char* buf = NULL
133136
cdef Py_ssize_t buf_len
134137
cdef char* cenc = NULL
135138
cdef char* cerr = NULL
136-
137-
PyObject_AsReadBuffer(packed, <const void**>&buf, &buf_len)
139+
cdef char buffer_supported = 0
140+
141+
if PyObject_CheckBuffer(packed):
142+
buffer_supported = 1
143+
if PyObject_GetBuffer(packed, &view, PyBUF_SIMPLE) == 0:
144+
if view.itemsize != 1:
145+
PyBuffer_Release(&view)
146+
raise ValueError("cannot unpack from multi-byte object")
147+
buf_len = view.len
148+
buf = <char*> view.buf
149+
else:
150+
PyObject_AsReadBuffer(packed, <const void**>&buf, &buf_len)
151+
PyErr_WarnEx(DeprecationWarning,
152+
"Unpacking %s requires old buffer protocol, "
153+
"which will be removed in msgpack 1.0." % type(packed),
154+
1)
138155

139156
if encoding is not None:
140157
if isinstance(encoding, unicode):
@@ -150,6 +167,8 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
150167
use_list, cenc, cerr,
151168
max_str_len, max_bin_len, max_array_len, max_map_len, max_ext_len)
152169
ret = unpack_construct(&ctx, buf, buf_len, &off)
170+
if buffer_supported:
171+
PyBuffer_Release(&view);
153172
if ret == 1:
154173
obj = unpack_data(&ctx)
155174
if off < buf_len:

‎msgpack/fallback.py‎

Lines changed: 43 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import sys
44
import array
55
import struct
6+
import warnings
67

78
if sys.version_info[0] == 3:
89
PY3 = True
@@ -246,16 +247,25 @@ def __init__(self, file_like=None, read_size=0, use_list=True,
246247
raise TypeError("`ext_hook` is not callable")
247248

248249
def feed(self, next_bytes):
249-
if isinstance(next_bytes, array.array):
250-
next_bytes = next_bytes.tostring()
251-
elif isinstance(next_bytes, bytearray):
252-
next_bytes = bytes(next_bytes)
253250
assert self._fb_feeding
254-
if (self._fb_buf_n + len(next_bytes) - self._fb_sloppiness
255-
> self._max_buffer_size):
251+
try:
252+
view = memoryview(next_bytes)
253+
except TypeError:
254+
# try to use legacy buffer protocol if 2.7, otherwise re-raise
255+
if not PY3:
256+
view = memoryview(buffer(next_bytes))
257+
warnings.warn("Unpacking %s requires old buffer protocol, "
258+
"which will be removed in msgpack 1.0." % type(next_bytes),
259+
DeprecationWarning)
260+
else:
261+
raise
262+
if view.itemsize != 1:
263+
raise ValueError("cannot unpack from multi-byte object")
264+
L = len(view)
265+
if self._fb_buf_n + L - self._fb_sloppiness > self._max_buffer_size:
256266
raise BufferFull
257-
self._fb_buf_n += len(next_bytes)
258-
self._fb_buffers.append(next_bytes)
267+
self._fb_buf_n += L
268+
self._fb_buffers.append(view)
259269

260270
def _fb_sloppy_consume(self):
261271
""" Gets rid of some of the used parts of the buffer. """
@@ -322,9 +332,10 @@ def _fb_read(self, n, write_bytes=None):
322332
return buffs[self._fb_buf_i][self._fb_buf_o - n:self._fb_buf_o]
323333

324334
# The remaining cases.
325-
ret = b''
326-
while len(ret) != n:
327-
sliced = n - len(ret)
335+
ret = []
336+
n_read = 0
337+
while n_read != n:
338+
sliced = n - n_read
328339
if self._fb_buf_i == len(buffs):
329340
if self._fb_feeding:
330341
break
@@ -334,34 +345,39 @@ def _fb_read(self, n, write_bytes=None):
334345
tmp = self.file_like.read(to_read)
335346
if not tmp:
336347
break
337-
buffs.append(tmp)
338-
self._fb_buf_n += len(tmp)
348+
tmpview = memoryview(tmp)
349+
assert tmpview.itemsize == 1
350+
buffs.append(tmpview)
351+
self._fb_buf_n += len(tmpview)
339352
continue
340-
ret += buffs[self._fb_buf_i][self._fb_buf_o:self._fb_buf_o + sliced]
353+
to_append = buffs[self._fb_buf_i][self._fb_buf_o:self._fb_buf_o + sliced]
354+
n_read += len(to_append)
355+
ret.append(to_append)
341356
self._fb_buf_o += sliced
342357
if self._fb_buf_o >= len(buffs[self._fb_buf_i]):
343358
self._fb_buf_o = 0
344359
self._fb_buf_i += 1
360+
ret = b''.join([view.tobytes() for view in ret])
345361
if len(ret) != n:
346362
self._fb_rollback()
347363
raise OutOfData
348364
if write_bytes is not None:
349365
write_bytes(ret)
350-
return ret
366+
return memoryview(ret)
351367

352368
def _read_header(self, execute=EX_CONSTRUCT, write_bytes=None):
353369
typ = TYPE_IMMEDIATE
354370
n = 0
355371
obj = None
356-
c = self._fb_read(1, write_bytes)
372+
c = self._fb_read(1, write_bytes).tobytes()
357373
b = ord(c)
358-
if b & 0b10000000 == 0:
374+
if b & 0b10000000 == 0:
359375
obj = b
360376
elif b & 0b11100000 == 0b11100000:
361377
obj = struct.unpack("b", c)[0]
362378
elif b & 0b11100000 == 0b10100000:
363379
n = b & 0b00011111
364-
obj = self._fb_read(n, write_bytes)
380+
obj = self._fb_read(n, write_bytes).tobytes()
365381
typ = TYPE_RAW
366382
if n > self._max_str_len:
367383
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
@@ -386,37 +402,37 @@ def _read_header(self, execute=EX_CONSTRUCT, write_bytes=None):
386402
n = struct.unpack("B", self._fb_read(1, write_bytes))[0]
387403
if n > self._max_bin_len:
388404
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
389-
obj = self._fb_read(n, write_bytes)
405+
obj = self._fb_read(n, write_bytes).tobytes()
390406
elif b == 0xc5:
391407
typ = TYPE_BIN
392408
n = struct.unpack(">H", self._fb_read(2, write_bytes))[0]
393409
if n > self._max_bin_len:
394410
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
395-
obj = self._fb_read(n, write_bytes)
411+
obj = self._fb_read(n, write_bytes).tobytes()
396412
elif b == 0xc6:
397413
typ = TYPE_BIN
398414
n = struct.unpack(">I", self._fb_read(4, write_bytes))[0]
399415
if n > self._max_bin_len:
400416
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
401-
obj = self._fb_read(n, write_bytes)
417+
obj = self._fb_read(n, write_bytes).tobytes()
402418
elif b == 0xc7: # ext 8
403419
typ = TYPE_EXT
404420
L, n = struct.unpack('Bb', self._fb_read(2, write_bytes))
405421
if L > self._max_ext_len:
406422
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
407-
obj = self._fb_read(L, write_bytes)
423+
obj = self._fb_read(L, write_bytes).tobytes()
408424
elif b == 0xc8: # ext 16
409425
typ = TYPE_EXT
410426
L, n = struct.unpack('>Hb', self._fb_read(3, write_bytes))
411427
if L > self._max_ext_len:
412428
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
413-
obj = self._fb_read(L, write_bytes)
429+
obj = self._fb_read(L, write_bytes).tobytes()
414430
elif b == 0xc9: # ext 32
415431
typ = TYPE_EXT
416432
L, n = struct.unpack('>Ib', self._fb_read(5, write_bytes))
417433
if L > self._max_ext_len:
418434
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
419-
obj = self._fb_read(L, write_bytes)
435+
obj = self._fb_read(L, write_bytes).tobytes()
420436
elif b == 0xca:
421437
obj = struct.unpack(">f", self._fb_read(4, write_bytes))[0]
422438
elif b == 0xcb:
@@ -467,19 +483,19 @@ def _read_header(self, execute=EX_CONSTRUCT, write_bytes=None):
467483
n = struct.unpack("B", self._fb_read(1, write_bytes))[0]
468484
if n > self._max_str_len:
469485
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
470-
obj = self._fb_read(n, write_bytes)
486+
obj = self._fb_read(n, write_bytes).tobytes()
471487
elif b == 0xda:
472488
typ = TYPE_RAW
473489
n = struct.unpack(">H", self._fb_read(2, write_bytes))[0]
474490
if n > self._max_str_len:
475491
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
476-
obj = self._fb_read(n, write_bytes)
492+
obj = self._fb_read(n, write_bytes).tobytes()
477493
elif b == 0xdb:
478494
typ = TYPE_RAW
479495
n = struct.unpack(">I", self._fb_read(4, write_bytes))[0]
480496
if n > self._max_str_len:
481497
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
482-
obj = self._fb_read(n, write_bytes)
498+
obj = self._fb_read(n, write_bytes).tobytes()
483499
elif b == 0xdc:
484500
n = struct.unpack(">H", self._fb_read(2, write_bytes))[0]
485501
if n > self._max_array_len:

‎test/test_buffer.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,3 +18,12 @@ def test_unpack_bytearray():
1818
assert [b'foo', b'bar'] == obj
1919
expected_type = bytes
2020
assert all(type(s) == expected_type for s in obj)
21+
22+
23+
def test_unpack_memoryview():
24+
buf = bytearray(packb(('foo', 'bar')))
25+
view = memoryview(buf)
26+
obj = unpackb(view, use_list=1)
27+
assert [b'foo', b'bar'] == obj
28+
expected_type = bytes
29+
assert all(type(s) == expected_type for s in obj)

0 commit comments

Comments
 (0)