Skip to content
Browse files

Add `autoreset` option to Packer.

  • Loading branch information...
1 parent 537a2ab commit 4adc6f194d9986eee4cfa4999d30862e4251d4cf @methane methane committed
Showing with 51 additions and 12 deletions.
  1. +36 −12 msgpack/_msgpack.pyx
  2. +15 −0 test/test_pack.py
View
48 msgpack/_msgpack.pyx
@@ -54,6 +54,16 @@ cdef class Packer(object):
packer = Packer()
astream.write(packer.pack(a))
astream.write(packer.pack(b))
+
+ Packer's constructor has some keyword arguments:
+
+ * *defaut* - Convert user type to builtin type that Packer supports.
+ See also simplejson's document.
+ * *encoding* - Convert unicode to bytes with this encoding. (default: 'utf-8')
+ * *unicode_erros* - Error handler for encoding unicode. (default: 'strict')
+ * *use_single_float* - Use single precision float type for float. (default: False)
+ * *autoreset* - Reset buffer after each pack and return it's content as `bytes`. (default: True).
+ If set this to false, use `bytes()` to get content and `.reset()` to clear buffer.
"""
cdef msgpack_packer pk
cdef object _default
@@ -62,6 +72,7 @@ cdef class Packer(object):
cdef char *encoding
cdef char *unicode_errors
cdef bool use_float
+ cdef bint autoreset
def __cinit__(self):
cdef int buf_size = 1024*1024
@@ -71,8 +82,9 @@ cdef class Packer(object):
self.pk.buf_size = buf_size
self.pk.length = 0
- def __init__(self, default=None, encoding='utf-8', unicode_errors='strict', use_single_float=False):
+ def __init__(self, default=None, encoding='utf-8', unicode_errors='strict', use_single_float=False, bint autoreset=1):
self.use_float = use_single_float
+ self.autoreset = autoreset
if default is not None:
if not PyCallable_Check(default):
raise TypeError("default must be a callable.")
@@ -182,9 +194,10 @@ cdef class Packer(object):
raise MemoryError
elif ret: # should not happen.
raise TypeError
- buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
- self.pk.length = 0
- return buf
+ if self.autoreset:
+ buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
+ self.pk.length = 0
+ return buf
def pack_array_header(self, size_t size):
cdef int ret = msgpack_pack_array(&self.pk, size)
@@ -192,9 +205,10 @@ cdef class Packer(object):
raise MemoryError
elif ret: # should not happen
raise TypeError
- buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
- self.pk.length = 0
- return buf
+ if self.autoreset:
+ buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
+ self.pk.length = 0
+ return buf
def pack_map_header(self, size_t size):
cdef int ret = msgpack_pack_map(&self.pk, size)
@@ -202,9 +216,10 @@ cdef class Packer(object):
raise MemoryError
elif ret: # should not happen
raise TypeError
- buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
- self.pk.length = 0
- return buf
+ if self.autoreset:
+ buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
+ self.pk.length = 0
+ return buf
def pack_map_pairs(self, object pairs):
"""
@@ -224,9 +239,18 @@ cdef class Packer(object):
raise MemoryError
elif ret: # should not happen
raise TypeError
- buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
+ if self.autoreset:
+ buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
+ self.pk.length = 0
+ return buf
+
+ def reset(self):
+ """Clear internal buffer."""
self.pk.length = 0
- return buf
+
+ def bytes(self):
+ """Return buffer content."""
+ return PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
def pack(object o, object stream, default=None, encoding='utf-8', unicode_errors='strict'):
View
15 test/test_pack.py
@@ -103,6 +103,21 @@ def testArraySize(sizes=[0, 5, 50, 1000]):
for size in sizes:
assert unpacker.unpack() == list(range(size))
+def test_manualreset(sizes=[0, 5, 50, 1000]):
+ packer = Packer(autoreset=False)
+ for size in sizes:
+ packer.pack_array_header(size)
+ for i in range(size):
+ packer.pack(i)
+
+ bio = six.BytesIO(packer.bytes())
+ unpacker = Unpacker(bio, use_list=1)
+ for size in sizes:
+ assert unpacker.unpack() == list(range(size))
+
+ packer.reset()
+ assert packer.bytes() == b''
+
def testMapSize(sizes=[0, 5, 50, 1000]):
bio = six.BytesIO()
packer = Packer()

0 comments on commit 4adc6f1

Please sign in to comment.
Something went wrong with that request. Please try again.