diff options
| -rw-r--r-- | msgpack/_msgpack.pyx | 37 | ||||
| -rw-r--r-- | test/test_unpack_raw.py | 32 |
2 files changed, 57 insertions, 12 deletions
diff --git a/msgpack/_msgpack.pyx b/msgpack/_msgpack.pyx index 7fcfcb0..71f9810 100644 --- a/msgpack/_msgpack.pyx +++ b/msgpack/_msgpack.pyx @@ -467,11 +467,16 @@ cdef class Unpacker(object): else: self.file_like = None - cdef object _unpack(self, execute_fn execute): + cdef object _unpack(self, execute_fn execute, object write_bytes): cdef int ret cdef object obj + cdef size_t prev_head while 1: + prev_head = self.buf_head ret = execute(&self.ctx, self.buf, self.buf_tail, &self.buf_head) + if write_bytes is not None: + write_bytes(PyBytes_FromStringAndSize(self.buf + prev_head, self.buf_head - prev_head)) + if ret == 1: obj = template_data(&self.ctx) template_init(&self.ctx) @@ -484,27 +489,35 @@ cdef class Unpacker(object): else: raise ValueError("Unpack failed: error = %d" % (ret,)) - def unpack(self): - """unpack one object""" - return self._unpack(template_construct) + def unpack(self, object write_bytes=None): + """ + unpack one object + + If write_bytes is not None, it will be called with parts of the raw message as it is unpacked. + """ + return self._unpack(template_construct, write_bytes) + + def skip(self, object write_bytes=None): + """ + read and ignore one object, returning None - def skip(self): - """read and ignore one object, returning None""" - return self._unpack(template_skip) + If write_bytes is not None, it will be called with parts of the raw message as it is unpacked. + """ + return self._unpack(template_skip, write_bytes) - def read_array_header(self): + def read_array_header(self, object write_bytes=None): """assuming the next object is an array, return its size n, such that the next n unpack() calls will iterate over its contents.""" - return self._unpack(read_array_header) + return self._unpack(read_array_header, write_bytes) - def read_map_header(self): + def read_map_header(self, object write_bytes=None): """assuming the next object is a map, return its size n, such that the next n * 2 unpack() calls will iterate over its key-value pairs.""" - return self._unpack(read_map_header) + return self._unpack(read_map_header, write_bytes) def __iter__(self): return self def __next__(self): - return self._unpack(template_construct) + return self._unpack(template_construct, None) # for debug. #def _buf(self): diff --git a/test/test_unpack_raw.py b/test/test_unpack_raw.py new file mode 100644 index 0000000..ed2e04b --- /dev/null +++ b/test/test_unpack_raw.py @@ -0,0 +1,32 @@ +"""Tests for cases where the user seeks to obtain packed msgpack objects""" + +from nose import main +from nose.tools import * +import six +from msgpack import Unpacker, packb + +def test_write_bytes(): + unpacker = Unpacker() + unpacker.feed(b'abc') + f = six.BytesIO() + assert_equal(unpacker.unpack(f.write), ord('a')) + assert_equal(f.getvalue(), b'a') + f.truncate(0) + assert_is_none(unpacker.skip(f.write)) + assert_equal(f.getvalue(), b'b') + f.truncate(0) + assert_is_none(unpacker.skip()) + assert_equal(f.getvalue(), b'') + +def test_write_bytes_multi_buffer(): + long_val = (5) * 100 + expected = packb(long_val) + unpacker = Unpacker(six.BytesIO(expected), read_size=3, max_buffer_size=3) + + f = six.BytesIO() + unpacked = unpacker.unpack(f.write) + assert_equal(unpacked, long_val) + assert_equal(f.getvalue(), expected) + +if __name__ == '__main__': + main() |