Merge pull request #161 from jfolz/feature/packbuffers

Support packing memoryview objects

2 files changed, 21 insertions, 2 deletions

diff --git a/msgpack/_packer.pyx b/msgpack/_packer.pyx
index c3ef1a4..b19d462 100644
--- a/msgpack/_packer.pyx
+++ b/msgpack/_packer.pyx
@@ -10,6 +10,11 @@ from msgpack.exceptions import PackValueError
 from msgpack import ExtType
 
 
+cdef extern from "Python.h":
+
+    int PyMemoryView_Check(object obj)
+
+
 cdef extern from "pack.h":
     struct msgpack_packer:
         char* buf
@@ -132,6 +137,7 @@ cdef class Packer(object):
         cdef size_t L
         cdef int default_used = 0
         cdef bint strict_types = self.strict_types
+        cdef Py_buffer view
 
         if nest_limit < 0:
             raise PackValueError("recursion limit exceeded.")
@@ -231,6 +237,17 @@ cdef class Packer(object):
                     for v in o:
                         ret = self._pack(v, nest_limit-1)
                         if ret != 0: break
+            elif PyMemoryView_Check(o):
+                if PyObject_GetBuffer(o, &view, PyBUF_SIMPLE) != 0:
+                    raise ValueError("could not get buffer for memoryview")
+                L = view.len
+                if L > (2**32)-1:
+                    PyBuffer_Release(&view);
+                    raise ValueError("memoryview is too large")
+                ret = msgpack_pack_bin(&self.pk, L)
+                if ret == 0:
+                    ret = msgpack_pack_raw_body(&self.pk, <char*>view.buf, L)
+                PyBuffer_Release(&view);
             elif not default_used and self._default:
                 o = self._default(o)
                 default_used = 1
diff --git a/msgpack/fallback.py b/msgpack/fallback.py
index 40c54a8..348e017 100644
--- a/msgpack/fallback.py
+++ b/msgpack/fallback.py
@@ -36,6 +36,8 @@ if hasattr(sys, 'pypy_version_info'):
             else:
                 self.builder = StringBuilder()
         def write(self, s):
+            if isinstance(s, memoryview):
+                s = s.tobytes()
             self.builder.append(s)
         def getvalue(self):
             return self.builder.build()
@@ -682,7 +684,7 @@ class Packer(object):
                     default_used = True
                     continue
                 raise PackValueError("Integer value out of range")
-            if self._use_bin_type and check(obj, bytes):
+            if self._use_bin_type and check(obj, (bytes, memoryview)):
                 n = len(obj)
                 if n <= 0xff:
                     self._buffer.write(struct.pack('>BB', 0xc4, n))
@@ -693,7 +695,7 @@ class Packer(object):
                 else:
                     raise PackValueError("Bytes is too large")
                 return self._buffer.write(obj)
-            if check(obj, (Unicode, bytes)):
+            if check(obj, (Unicode, bytes, memoryview)):
                 if check(obj, Unicode):
                     if self._encoding is None:
                         raise TypeError(