diff options
| -rw-r--r-- | msgpack/_msgpack.pyx | 39 | ||||
| -rw-r--r-- | test/test_pack.py | 11 |
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() |