Skip to content

Commit 537a2ab

Browse files
committed
Add Packer.pack_pairs.
1 parent 0c7ab7c commit 537a2ab

File tree

2 files changed

+44
-6
lines changed

2 files changed

+44
-6
lines changed

msgpack/_msgpack.pyx

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -186,18 +186,49 @@ cdef class Packer(object):
186186
self.pk.length = 0
187187
return buf
188188

189-
cpdef pack_array_header(self, size_t size):
190-
msgpack_pack_array(&self.pk, size)
189+
def pack_array_header(self, size_t size):
190+
cdef int ret = msgpack_pack_array(&self.pk, size)
191+
if ret == -1:
192+
raise MemoryError
193+
elif ret: # should not happen
194+
raise TypeError
191195
buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
192196
self.pk.length = 0
193197
return buf
194198

195-
cpdef pack_map_header(self, size_t size):
196-
msgpack_pack_map(&self.pk, size)
199+
def pack_map_header(self, size_t size):
200+
cdef int ret = msgpack_pack_map(&self.pk, size)
201+
if ret == -1:
202+
raise MemoryError
203+
elif ret: # should not happen
204+
raise TypeError
197205
buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
198206
self.pk.length = 0
199207
return buf
200208

209+
def pack_map_pairs(self, object pairs):
210+
"""
211+
Pack *pairs* as msgpack map type.
212+
213+
*pairs* should sequence of pair.
214+
(`len(pairs)` and `for k, v in *pairs*:` should be supported.)
215+
"""
216+
cdef int ret = msgpack_pack_map(&self.pk, len(pairs))
217+
if ret == 0:
218+
for k, v in pairs:
219+
ret = self._pack(k)
220+
if ret != 0: break
221+
ret = self._pack(v)
222+
if ret != 0: break
223+
if ret == -1:
224+
raise MemoryError
225+
elif ret: # should not happen
226+
raise TypeError
227+
buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
228+
self.pk.length = 0
229+
return buf
230+
231+
201232
def pack(object o, object stream, default=None, encoding='utf-8', unicode_errors='strict'):
202233
"""
203234
pack an object `o` and write it to stream)."""

test/test_pack.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -118,8 +118,6 @@ def testMapSize(sizes=[0, 5, 50, 1000]):
118118
assert unpacker.unpack() == dict((i, i * 2) for i in range(size))
119119

120120

121-
122-
123121
class odict(dict):
124122
'''Reimplement OrderedDict to run test on Python 2.6'''
125123
def __init__(self, seq):
@@ -144,5 +142,14 @@ def pair_hook(seq):
144142
assert_equal(unpackb(packb(od), object_pairs_hook=pair_hook, use_list=1), seq)
145143

146144

145+
def test_pairlist():
146+
pairlist = [(b'a', 1), (2, b'b'), (b'foo', b'bar')]
147+
packer = Packer()
148+
packed = packer.pack_map_pairs(pairlist)
149+
unpacked = unpackb(packed, object_pairs_hook=list)
150+
assert pairlist == unpacked
151+
152+
153+
147154
if __name__ == '__main__':
148155
main()

0 commit comments

Comments
 (0)