Merge pull request #23 from jnothman/write_bytes

Allow packed data to be captured while executing skip(), etc.

Author: methane

msgpack/_msgpack.pyx

@@ -479,11 +479,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)
@@ -496,27 +501,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):

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()