Add `Packer.pack_pairs`.

2 files changed, 44 insertions, 6 deletions

diff --git a/msgpack/_msgpack.pyx b/msgpack/_msgpack.pyx
index 52d63cc..2d97f21 100644
--- a/msgpack/_msgpack.pyx
+++ b/msgpack/_msgpack.pyx
@@ -186,18 +186,49 @@ cdef class Packer(object):
         self.pk.length = 0
         return buf
 
-    cpdef pack_array_header(self, size_t size):
-        msgpack_pack_array(&self.pk, size)
+    def pack_array_header(self, size_t size):
+        cdef int ret = msgpack_pack_array(&self.pk, size)
+        if ret == -1:
+            raise MemoryError
+        elif ret:  # should not happen
+            raise TypeError
         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)
+    def pack_map_header(self, size_t size):
+        cdef int ret = msgpack_pack_map(&self.pk, size)
+        if ret == -1:
+            raise MemoryError
+        elif ret:  # should not happen
+            raise TypeError
         buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
         self.pk.length = 0
         return buf
 
+    def pack_map_pairs(self, object pairs):
+        """
+        Pack *pairs* as msgpack map type.
+
+        *pairs* should sequence of pair.
+        (`len(pairs)` and `for k, v in *pairs*:` should be supported.)
+        """
+        cdef int ret = msgpack_pack_map(&self.pk, len(pairs))
+        if ret == 0:
+            for k, v in pairs:
+                ret = self._pack(k)
+                if ret != 0: break
+                ret = self._pack(v)
+                if ret != 0: break
+        if ret == -1:
+            raise MemoryError
+        elif ret:  # should not happen
+            raise TypeError
+        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'):
     """
     pack an object `o` and write it to stream)."""
diff --git a/test/test_pack.py b/test/test_pack.py
index 21c2bd7..b918f8e 100644
--- a/test/test_pack.py
+++ b/test/test_pack.py
@@ -118,8 +118,6 @@ def testMapSize(sizes=[0, 5, 50, 1000]):
         assert unpacker.unpack() == dict((i, i * 2) for i in range(size))
 
 
-
-
 class odict(dict):
     '''Reimplement OrderedDict to run test on Python 2.6'''
     def __init__(self, seq):
@@ -144,5 +142,14 @@ def test_odict():
     assert_equal(unpackb(packb(od), object_pairs_hook=pair_hook, use_list=1), seq)
 
 
+def test_pairlist():
+    pairlist = [(b'a', 1), (2, b'b'), (b'foo', b'bar')]
+    packer = Packer()
+    packed = packer.pack_map_pairs(pairlist)
+    unpacked = unpackb(packed, object_pairs_hook=list)
+    assert pairlist == unpacked
+
+
+
 if __name__ == '__main__':
     main()