Packer.pack_array/map_header to correspond to read functions

2 files changed, 40 insertions, 0 deletions

diff --git a/msgpack/_msgpack.pyx b/msgpack/_msgpack.pyx
index 7131d1f..18a75ca 100644
--- a/msgpack/_msgpack.pyx
+++ b/msgpack/_msgpack.pyx
@@ -178,6 +178,17 @@ cdef class Packer(object):
         self.pk.length = 0
         return buf
 
+    cpdef pack_array_header(self, size_t size):
+        msgpack_pack_array(&self.pk, size)
+        buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
+        self.pk.length = 0
+        return buf
+
+    cpdef pack_map_header(self, size_t size):
+        msgpack_pack_map(&self.pk, size)
+        buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
+        self.pk.length = 0
+        return buf
 
 def pack(object o, object stream, default=None, encoding='utf-8', unicode_errors='strict'):
     """
diff --git a/test/test_pack.py b/test/test_pack.py
index b216c46..937141d 100644
--- a/test/test_pack.py
+++ b/test/test_pack.py
@@ -91,6 +91,35 @@ def testPackFloat():
     assert_equal(packb(1.0, use_single_float=True),  b'\xca' + struct.pack('>f', 1.0))
     assert_equal(packb(1.0, use_single_float=False), b'\xcb' + struct.pack('>d', 1.0))
 
+def testArraySize(sizes=[0, 5, 50, 1000]):
+    bio = six.BytesIO()
+    packer = Packer()
+    for size in sizes:
+        bio.write(packer.pack_array_header(size))
+        for i in range(size):
+            bio.write(packer.pack(i))
+
+    bio.seek(0)
+    unpacker = Unpacker(bio)
+    for size in sizes:
+        assert unpacker.unpack() == tuple(range(size))
+
+def testMapSize(sizes=[0, 5, 50, 1000]):
+    bio = six.BytesIO()
+    packer = Packer()
+    for size in sizes:
+        bio.write(packer.pack_map_header(size))
+        for i in range(size):
+            bio.write(packer.pack(i)) # key
+            bio.write(packer.pack(i * 2)) # value
+
+    bio.seek(0)
+    unpacker = Unpacker(bio)
+    for size in sizes:
+        assert unpacker.unpack() == {i: i * 2 for i in range(size)}
+
+
+
 
 class odict(dict):
     '''Reimplement OrderedDict to run test on Python 2.6'''