8000 Support packing memoryview objects · sugarguo/msgpack-python@31adc5a · GitHub
[go: up one dir, main page]
Skip to content

Commit 31adc5a

Browse files
jfolzjfolz
committed
Support packing memoryview objects
1 parent 8036cb4 commit 31adc5a

File tree

3 files changed

+33
-2
lines changed

3 files changed

+33
-2
lines changed

‎msgpack/_packer.pyx

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,11 @@ from msgpack.exceptions import PackValueError
1010
from msgpack import ExtType
1111

1212

13+
cdef extern from "Python.h":
14+
15+
int PyMemoryView_Check(object obj)
16+
17+
1318
cdef extern from "pack.h":
1419
struct msgpack_packer:
1520
char* buf
@@ -132,6 +137,7 @@ cdef class Packer(object):
132137
cdef size_t L
133138
cdef int default_used = 0
134139
cdef bint strict_types = self.strict_types
140+
cdef Py_buffer view
135141

136142
if nest_limit < 0:
137143
raise PackValueError("recursion limit exceeded.")
@@ -231,6 +237,17 @@ cdef class Packer(object):
231237
for v in o:
232238
ret = self._pack(v, nest_limit-1)
233239
if ret != 0: break
240+
elif PyMemoryView_Check(o):
241+
if PyObject_GetBuffer(o, &view, PyBUF_SIMPLE) != 0:
242+
raise ValueError("could not get buffer for memoryview")
243+
L = view.len
244+
if L > (2**32)-1:
245+
PyBuffer_Release(&view);
246+
raise ValueError("memoryview is too large")
247+
ret = msgpack_pack_bin(&self.pk, L)
248+
if ret == 0:
249+
ret = msgpack_pack_raw_body(&self.pk, <char*>view.buf, L)
250+
PyBuffer_Release(&view);
234251
elif not default_used and self._default:
235252
o = self._default(o)
236253
default_used = 1

‎msgpack/fallback.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ def __init__(self, s=b''):
3636
else:
3737
self.builder = StringBuilder()
3838
def write(self, s):
39+
if isinstance(s, memoryview):
40+
s = s.tobytes()
3941
self.builder.append(s)
4042
def getvalue(self):
4143
return self.builder.build()
@@ -682,7 +684,7 @@ def _pack(self, obj, nest_limit=DEFAULT_RECURSE_LIMIT,
682684
default_used = True
683685
continue
684686
raise PackValueError("Integer value out of range")
685-
if self._use_bin_type and check(obj, bytes):
687+
if self._use_bin_type and check(obj, (bytes, memoryview)):
686688
n = len(obj)
687689
if n <= 0xff:
688690
self._buffer.write(struct.pack('>BB', 0xc4, n))
@@ -693,7 +695,7 @@ def _pack(self, obj, nest_limit=DEFAULT_RECURSE_LIMIT,
693695
else:
694696
raise PackValueError("Bytes is too large")
695697
return self._buffer.write(obj)
696-
if check(obj, (Unicode, bytes)):
698+
if check(obj, (Unicode, bytes, memoryview)):
697699
if check(obj, Unicode):
698700
if self._encoding is None:
699701
raise TypeError(

‎test/test_memoryview.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
#!/usr/bin/env python
2+
# coding: utf-8
3+
4+
5+
from msgpack import packb, unpackb
6+
7+
8+
def test_pack_memoryview():
9+
data = bytearray(range(256))
10+
view = memoryview(data)
11+
unpacked = unpackb(packb(view))
12+
assert data == unpacked

0 commit comments

Comments
 (0)
0