8000 enable unpacking from memoryview · msgpack/msgpack-python@31f76fb · GitHub
[go: up one dir, main page]
Skip to content

Commit 31f76fb

Browse files
jfolzjfolz
committed
enable unpacking from memoryview
1 parent b78c0c5 commit 31f76fb

File tree

3 files changed

+91
-15
lines changed

3 files changed

+91
-15
lines changed

‎msgpack/_unpacker.pyx

Lines changed: 58 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,16 +8,24 @@ from cpython.bytes cimport (
88
)
99
from cpython.buffer cimport (
1010
Py_buffer,
11-
PyBuffer_Release,
11+
PyObject_CheckBuffer,
1212
PyObject_GetBuffer,
13+
PyBuffer_Release,
14+
PyBuffer_IsContiguous,
15+
PyBuffer_FillInfo,
16+
PyBUF_READ,
1317
PyBUF_SIMPLE,
18+
PyBUF_FULL_RO,
1419
)
1520
from cpython.mem cimport PyMem_Malloc, PyMem_Free
1621
from cpython.object cimport PyCallable_Check
22+
from cpython.ref cimport Py_DECREF
23+
from cpython.exc cimport PyErr_WarnEx
1724

1825
cdef extern from "Python.h":
1926
ctypedef struct PyObject
2027
cdef int PyObject_AsReadBuffer(object o, const void** buff, Py_ssize_t* buf_len) except -1
28+
object PyMemoryView_GetContiguous(object obj, int buffertype, char order)
2129

2230
from libc.stdlib cimport *
2331
from libc.string cimport *
@@ -110,6 +118,40 @@ cdef inline init_ctx(unpack_context *ctx,
110118
def default_read_extended_type(typecode, data):
111119
raise NotImplementedError("Cannot decode extended type with typecode=%d" % typecode)
112120

121+
cdef inline int get_data_from_buffer(object obj,
122+
Py_buffer *view,
123+
char **buf,
124+
Py_ssize_t *buffer_len,
125+
int *new_protocol) except 0:
126+
cdef object contiguous
127+
cdef Py_buffer tmp
128+
if PyObject_CheckBuffer(obj):
129+
new_protocol[0] = 1
130+
if PyObject_GetBuffer(obj, view, PyBUF_FULL_RO) == -1:
131+
raise
132+
if view.itemsize != 1:
133+
PyBuffer_Release(view)
134+
raise BufferError("cannot unpack from multi-byte object")
135+
if PyBuffer_IsContiguous(view, 'A') == 0:
136+
contiguous = PyMemoryView_GetContiguous(obj, PyBUF_READ, 'C')
137+
PyObject_GetBuffer(contiguous, &tmp, PyBUF_SIMPLE)
138+
PyBuffer_Release(view)
139+
PyBuffer_FillInfo(view, contiguous, tmp.buf, tmp.len, 1, 0)
140+
Py_DECREF(contiguous)
141+
buffer_len[0] = view.len
142+
buf[0] = <char*> view.buf
143+
return 1
144+
else:
145+
new_protocol[0] = 0
146+
if PyObject_AsReadBuffer(obj, <const void**> buf, buffer_len) == -1:
147+
raise BufferError("could not get memoryview")
148+
PyErr_WarnEx(RuntimeWarning,
149+
"using old buffer interface to unpack %s; "
150+
"this leads to unpacking errors if slicing is used and "
151+
"will be removed in a future version" % type(obj),
152+
1)
153+
return 1
154+
113155
def unpackb(object packed, object object_hook=None, object list_hook=None,
114156
bint use_list= 8000 1, encoding=None, unicode_errors="strict",
115157
object_pairs_hook=None, ext_hook=ExtType,
@@ -129,12 +171,14 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
129171
cdef Py_ssize_t off = 0
130172
cdef int ret
131173

132-
cdef char* buf
174+
cdef Py_buffer view
175+
cdef char* buf = NULL
133176
cdef Py_ssize_t buf_len
134177
cdef char* cenc = NULL
135178
cdef char* cerr = NULL
179+
cdef int new_protocol = 0
136180

137-
PyObject_AsReadBuffer(packed, <const void**>&buf, &buf_len)
181+
get_data_from_buffer(packed, &view, &buf, &buf_len, &new_protocol)
138182

139183
if encoding is not None:
140184
if isinstance(encoding, unicode):
@@ -150,6 +194,8 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
150194
use_list, cenc, cerr,
151195
max_str_len, max_bin_len, max_array_len, max_map_len, max_ext_len)
152196
ret = unpack_construct(&ctx, buf, buf_len, &off)
197+
if new_protocol:
198+
PyBuffer_Release(&view);
153199
if ret == 1:
154200
obj = unpack_data(&ctx)
155201
if off < buf_len:
@@ -335,14 +381,20 @@ cdef class Unpacker(object):
335381
def feed(self, object next_bytes):
336382
"""Append `next_bytes` to internal buffer."""
337383
cdef Py_buffer pybuff
384+
cdef int new_protocol = 0
385+
cdef char* buf
386+
cdef Py_ssize_t buf_len
387+
338388
if self.file_like is not None:
339389
raise AssertionError(
340390
"unpacker.feed() is not be able to use with `file_like`.")
341-
PyObject_GetBuffer(next_bytes, &pybuff, PyBUF_SIMPLE)
391+
392+
get_data_from_buffer(next_bytes, &pybuff, &buf, &buf_len, &new_protocol)
342393
try:
< 9E88 /td>
343-
self.append_buffer(<char*>pybuff.buf, pybuff.len)
394+
self.append_buffer(buf, buf_len)
344395
finally:
345-
PyBuffer_Release(&pybuff)
396+
if new_protocol:
397+
PyBuffer_Release(&pybuff)
346398

347399
cdef append_buffer(self, void* _buf, Py_ssize_t _buf_len):
348400
cdef:

‎msgpack/fallback.py

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
"""Fallback pure Python implementation of msgpack"""
22

33
import sys
4-
import array
54
import struct
5+
import warnings
66

77
if sys.version_info[0] == 3:
88
PY3 = True
@@ -46,6 +46,7 @@ def getvalue(self):
4646
from io import BytesIO as StringIO
4747
newlist_hint = lambda size: []
4848

49+
4950
from msgpack.exceptions import (
5051
BufferFull,
5152
OutOfData,
@@ -79,6 +80,24 @@ def _check_type_strict(obj, t, type=type, tuple=tuple):
7980
return type(obj) is t
8081

8182

83+
def _get_data_from_buffer(obj):
84+
try:
85+
view = memoryview(obj)
86+
except TypeError:
87+
# try to use legacy buffer protocol if 2.7, otherwise re-raise
88+
if not PY3:
89+
view = memoryview(buffer(obj))
90+
warnings.warn("using old buffer interface to unpack %s; "
91+
"this leads to unpacking errors if slicing is used and "
92+
"will be removed in a future version" % type(obj),
93+
RuntimeWarning)
94+
else:
95+
raise
96+
if view.itemsize != 1:
97+
raise ValueError("cannot unpack from multi-byte object")
98+
return view
99+
100+
82101
def unpack(stream, **kwargs):
83102
"""
84103
Unpack an object from `stream`.
@@ -239,17 +258,13 @@ def __init__(self, file_like=None, read_size=0, use_list=True,
239258
raise TypeError("`ext_hook` is not callable")
240259

241260
def feed(self, next_bytes):
242-
if isinstance(next_bytes, array.array):
243-
next_bytes = next_bytes.tostring()
244-
if not isinstance(next_bytes, (bytes, bytearray)):
245-
raise TypeError("next_bytes should be bytes, bytearray or array.array")
246261
assert self._feeding
247-
248-
if (len(self._buffer) - self._buff_i + len(next_bytes) > self._max_buffer_size):
262+
view = _get_data_from_buffer(next_bytes)
263+
if (len(self._buffer) - self._buff_i + len(view) > self._max_buffer_size):
249264
raise BufferFull
250265
# bytes + bytearray -> bytearray
251266
# So cast before append
252-
self._buffer += bytes(next_bytes)
267+
self._buffer += view.tobytes()
253268

254269
def _consume(self):
255270
""" Gets rid of the used parts of the buffer. """
@@ -316,7 +331,7 @@ def _read_header(self, execute=EX_CONSTRUCT):
316331
obj = None
317332
c = self._read(1)
318333
b = ord(c)
319-
if b & 0b10000000 == 0:
334+
if b & 0b10000000 == 0:
320335
obj = b
321336
elif b & 0b11100000 == 0b11100000:
322337
obj = struct.unpack("b", c)[0]

‎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)
0