Pure Python packer supports memoryview of multi byte items.

3 files changed, 159 insertions, 41 deletions

diff --git a/msgpack/_packer.pyx b/msgpack/_packer.pyx
index b1a912b..e07b194 100644
--- a/msgpack/_packer.pyx
+++ b/msgpack/_packer.pyx
@@ -39,6 +39,7 @@ cdef extern from "pack.h":
     int msgpack_pack_ext(msgpack_packer* pk, char typecode, size_t l)
 
 cdef int DEFAULT_RECURSE_LIMIT=511
+cdef size_t ITEM_LIMIT = (2**32)-1
 
 
 cdef class Packer(object):
@@ -178,7 +179,7 @@ cdef class Packer(object):
                    ret = msgpack_pack_double(&self.pk, dval)
             elif PyBytes_CheckExact(o) if strict_types else PyBytes_Check(o):
                 L = len(o)
-                if L > (2**32)-1:
+                if L > ITEM_LIMIT:
                     raise PackValueError("bytes is too large")
                 rawval = o
                 ret = msgpack_pack_bin(&self.pk, L)
@@ -189,7 +190,7 @@ cdef class Packer(object):
                     raise TypeError("Can't encode unicode string: no encoding is specified")
                 o = PyUnicode_AsEncodedString(o, self.encoding, self.unicode_errors)
                 L = len(o)
-                if L > (2**32)-1:
+                if L > ITEM_LIMIT:
                     raise PackValueError("unicode string is too large")
                 rawval = o
                 ret = msgpack_pack_raw(&self.pk, L)
@@ -198,7 +199,7 @@ cdef class Packer(object):
             elif PyDict_CheckExact(o):
                 d = <dict>o
                 L = len(d)
-                if L > (2**32)-1:
+                if L > ITEM_LIMIT:
                     raise PackValueError("dict is too large")
                 ret = msgpack_pack_map(&self.pk, L)
                 if ret == 0:
@@ -209,7 +210,7 @@ cdef class Packer(object):
                         if ret != 0: break
             elif not strict_types and PyDict_Check(o):
                 L = len(o)
-                if L > (2**32)-1:
+                if L > ITEM_LIMIT:
                     raise PackValueError("dict is too large")
                 ret = msgpack_pack_map(&self.pk, L)
                 if ret == 0:
@@ -223,13 +224,13 @@ cdef class Packer(object):
                 longval = o.code
                 rawval = o.data
                 L = len(o.data)
-                if L > (2**32)-1:
+                if L > ITEM_LIMIT:
                     raise PackValueError("EXT data is too large")
                 ret = msgpack_pack_ext(&self.pk, longval, L)
                 ret = msgpack_pack_raw_body(&self.pk, rawval, L)
             elif PyList_CheckExact(o) if strict_types else (PyTuple_Check(o) or PyList_Check(o)):
                 L = len(o)
-                if L > (2**32)-1:
+                if L > ITEM_LIMIT:
                     raise PackValueError("list is too large")
                 ret = msgpack_pack_array(&self.pk, L)
                 if ret == 0:
@@ -240,7 +241,7 @@ cdef class Packer(object):
                 if PyObject_GetBuffer(o, &view, PyBUF_SIMPLE) != 0:
                     raise PackValueError("could not get buffer for memoryview")
                 L = view.len
-                if L > (2**32)-1:
+                if L > ITEM_LIMIT:
                     PyBuffer_Release(&view);
                     raise PackValueError("memoryview is too large")
                 ret = msgpack_pack_bin(&self.pk, L)
@@ -271,8 +272,8 @@ cdef class Packer(object):
         msgpack_pack_ext(&self.pk, typecode, len(data))
         msgpack_pack_raw_body(&self.pk, data, len(data))
 
-    def pack_array_header(self, long long size):
-        if size > (2**32-1):
+    def pack_array_header(self, size_t size):
+        if size > ITEM_LIMIT:
             raise PackValueError
         cdef int ret = msgpack_pack_array(&self.pk, size)
         if ret == -1:
@@ -284,8 +285,8 @@ cdef class Packer(object):
             self.pk.length = 0
             return buf
 
-    def pack_map_header(self, long long size):
-        if size > (2**32-1):
+    def pack_map_header(self, size_t size):
+        if size > ITEM_LIMIT:
             raise PackValueError
         cdef int ret = msgpack_pack_map(&self.pk, size)
         if ret == -1:
diff --git a/msgpack/fallback.py b/msgpack/fallback.py
index d8c5d73..abed3d9 100644
--- a/msgpack/fallback.py
+++ b/msgpack/fallback.py
@@ -685,35 +685,28 @@ class Packer(object):
                     default_used = True
                     continue
                 raise PackOverflowError("Integer value out of range")
-            if self._use_bin_type and check(obj, (bytes, memoryview)):
+            if check(obj, bytes):
                 n = len(obj)
-                if n <= 0xff:
-                    self._buffer.write(struct.pack('>BB', 0xc4, n))
-                elif n <= 0xffff:
-                    self._buffer.write(struct.pack(">BH", 0xc5, n))
-                elif n <= 0xffffffff:
-                    self._buffer.write(struct.pack(">BI", 0xc6, n))
-                else:
+                if n >= 2**32:
                     raise PackValueError("Bytes is too large")
+                self._fb_pack_bin_header(n)
                 return self._buffer.write(obj)
-            if check(obj, (Unicode, bytes, memoryview)):
-                if check(obj, Unicode):
-                    if self._encoding is None:
-                        raise TypeError(
-                            "Can't encode unicode string: "
-                            "no encoding is specified")
-                    obj = obj.encode(self._encoding, self._unicode_errors)
+            if check(obj, Unicode):
+                if self._encoding is None:
+                    raise TypeError(
+                        "Can't encode unicode string: "
+                        "no encoding is specified")
+                obj = obj.encode(self._encoding, self._unicode_errors)
                 n = len(obj)
-                if n <= 0x1f:
-                    self._buffer.write(struct.pack('B', 0xa0 + n))
-                elif self._use_bin_type and n <= 0xff:
-                    self._buffer.write(struct.pack('>BB', 0xd9, n))
-                elif n <= 0xffff:
-                    self._buffer.write(struct.pack(">BH", 0xda, n))
-                elif n <= 0xffffffff:
-                    self._buffer.write(struct.pack(">BI", 0xdb, n))
-                else:
+                if n >= 2**32:
                     raise PackValueError("String is too large")
+                self._fb_pack_raw_header(n)
+                return self._buffer.write(obj)
+            if check(obj, memoryview):
+                n = len(obj) * obj.itemsize
+                if n >= 2**32:
+                    raise PackValueError("Memoryview is too large")
+                self._fb_pack_bin_header(n)
                 return self._buffer.write(obj)
             if check(obj, float):
                 if self._use_float:
@@ -852,6 +845,30 @@ class Packer(object):
             self._pack(k, nest_limit - 1)
             self._pack(v, nest_limit - 1)
 
+    def _fb_pack_raw_header(self, n):
+        if n <= 0x1f:
+            self._buffer.write(struct.pack('B', 0xa0 + n))
+        elif self._use_bin_type and n <= 0xff:
+            self._buffer.write(struct.pack('>BB', 0xd9, n))
+        elif n <= 0xffff:
+            self._buffer.write(struct.pack(">BH", 0xda, n))
+        elif n <= 0xffffffff:
+            self._buffer.write(struct.pack(">BI", 0xdb, n))
+        else:
+            raise PackValueError('Raw is too large')
+
+    def _fb_pack_bin_header(self, n):
+        if not self._use_bin_type:
+            return self._fb_pack_raw_header(n)
+        elif n <= 0xff:
+            return self._buffer.write(struct.pack('>BB', 0xc4, n))
+        elif n <= 0xffff:
+            return self._buffer.write(struct.pack(">BH", 0xc5, n))
+        elif n <= 0xffffffff:
+            return self._buffer.write(struct.pack(">BI", 0xc6, n))
+        else:
+            raise PackValueError('Bin is too large')
+
     def bytes(self):
         return self._buffer.getvalue()
 
diff --git a/test/test_memoryview.py b/test/test_memoryview.py
index aed5069..f6d74ed 100644
--- a/test/test_memoryview.py
+++ b/test/test_memoryview.py
@@ -1,12 +1,112 @@
 #!/usr/bin/env python
 # coding: utf-8
 
-
+from array import array
 from msgpack import packb, unpackb
+import sys
+
+
+# For Python < 3:
+#  - array type only supports old buffer interface
+#  - array.frombytes is not available, must use deprecated array.fromstring
+if sys.version_info[0] < 3:
+    def make_memoryview(obj):
+        return memoryview(buffer(obj))
+
+    def make_array(f, data):
+        a = array(f)
+        a.fromstring(data)
+        return a
+
+    def get_data(a):
+        return a.tostring()
+else:
+    make_memoryview = memoryview
+
+    def make_array(f, data):
+        a = array(f)
+        a.frombytes(data)
+        return a
+
+    def get_data(a):
+        return a.tobytes()
+
+
+def _runtest(format, nbytes, expected_header, expected_prefix, use_bin_type):
+    # create a new array
+    original_array = array(format)
+    original_array.fromlist([255] * (nbytes // original_array.itemsize))
+    original_data = get_data(original_array)
+    view = make_memoryview(original_array)
+
+    # pack, unpack, and reconstruct array
+    packed = packb(view, use_bin_type=use_bin_type)
+    unpacked = unpackb(packed)
+    reconstructed_array = make_array(format, unpacked)
+
+    # check that we got the right amount of data
+    assert len(original_data) == nbytes
+    # check packed header
+    assert packed[:1] == expected_header
+    # check packed length prefix, if any
+    assert packed[1:1+len(expected_prefix)] == expected_prefix
+    # check packed data
+    assert packed[1+len(expected_prefix):] == original_data
+    # check array unpacked correctly
+    assert original_array == reconstructed_array
+
+
+def test_fixstr_from_byte():
+    _runtest('B', 1, b'\xa1', b'', False)
+    _runtest('B', 31, b'\xbf', b'', False)
+
+
+def test_fixstr_from_float():
+    _runtest('f', 4, b'\xa4', b'', False)
+    _runtest('f', 28, b'\xbc', b'', False)
+
+
+def test_str16_from_byte():
+    _runtest('B', 2**8, b'\xda', b'\x01\x00', False)
+    _runtest('B', 2**16-1, b'\xda', b'\xff\xff', False)
+
+
+def test_str16_from_float():
+    _runtest('f', 2**8, b'\xda', b'\x01\x00', False)
+    _runtest('f', 2**16-4, b'\xda', b'\xff\xfc', False)
+
+
+def test_str32_from_byte():
+    _runtest('B', 2**16, b'\xdb', b'\x00\x01\x00\x00', False)
+
+
+def test_str32_from_float():
+    _runtest('f', 2**16, b'\xdb', b'\x00\x01\x00\x00', False)
+
+
+def test_bin8_from_byte():
+    _runtest('B', 1, b'\xc4', b'\x01', True)
+    _runtest('B', 2**8-1, b'\xc4', b'\xff', True)
+
+
+def test_bin8_from_float():
+    _runtest('f', 4, b'\xc4', b'\x04', True)
+    _runtest('f', 2**8-4, b'\xc4', b'\xfc', True)
+
+
+def test_bin16_from_byte():
+    _runtest('B', 2**8, b'\xc5', b'\x01\x00', True)
+    _runtest('B', 2**16-1, b'\xc5', b'\xff\xff', True)
+
+
+def test_bin16_from_float():
+    _runtest('f', 2**8, b'\xc5', b'\x01\x00', True)
+    _runtest('f', 2**16-4, b'\xc5', b'\xff\xfc', True)
+
+
+def test_bin32_from_byte():
+    _runtest('B', 2**16, b'\xc6', b'\x00\x01\x00\x00', True)
 
 
-def test_pack_memoryview():
-    data = bytearray(range(256))
-    view = memoryview(data)
-    unpacked = unpackb(packb(view))
-    assert data == unpacked
+def test_bin32_from_float():
+    _runtest('f', 2**16, b'\xc6', b'\x00\x01\x00\x00', True)