read_array/map_header functionality

3 files changed, 139 insertions, 0 deletions

diff --git a/msgpack/_msgpack.pyx b/msgpack/_msgpack.pyx
index 0fc3739..7131d1f 100644
--- a/msgpack/_msgpack.pyx
+++ b/msgpack/_msgpack.pyx
@@ -212,6 +212,8 @@ cdef extern from "unpack.h":
                                size_t len, size_t* off) except -1
     execute_fn template_construct
     execute_fn template_skip
+    execute_fn read_array_header
+    execute_fn read_map_header
     void template_init(template_context* ctx)
     object template_data(template_context* ctx)
 
@@ -482,6 +484,14 @@ cdef class Unpacker(object):
         """read and ignore one object, returning None"""
         return self._unpack(template_skip)
 
+    def read_array_header(self):
+        """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)
+
+    def read_map_header(self):
+        """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)
+
     def __iter__(self):
         return self
 
diff --git a/msgpack/unpack_template.h b/msgpack/unpack_template.h
index e0cf42e..69ef6e2 100644
--- a/msgpack/unpack_template.h
+++ b/msgpack/unpack_template.h
@@ -408,6 +408,10 @@ _end:
 #undef construct_cb
 }
 
+#undef SWITCH_RANGE_BEGIN
+#undef SWITCH_RANGE
+#undef SWITCH_RANGE_DEFAULT
+#undef SWITCH_RANGE_END
 #undef push_simple_value
 #undef push_fixed_value
 #undef push_variable_value
@@ -415,8 +419,67 @@ _end:
 #undef again_fixed_trail_if_zero
 #undef start_container
 
+template <unsigned int fixed_offset, unsigned int var_offset>
+msgpack_unpack_func(int, _container_header)(msgpack_unpack_struct(_context)* ctx, const char* data, size_t len, size_t* off)
+{
+	assert(len >= *off);
+	uint32_t size;
+	const unsigned char *const p = (unsigned char*)data + *off;
+
+#define inc_offset(inc) \
+	if (len - *off < inc) \
+		return 0; \
+	*off += inc;
+
+	switch (*p) {
+	case var_offset:
+		inc_offset(3);
+		size = _msgpack_load16(uint16_t, p + 1);
+		break;
+	case var_offset + 1:
+		inc_offset(5);
+		size = _msgpack_load32(uint32_t, p + 1);
+		break;
+#ifdef USE_CASE_RANGE
+	case fixed_offset + 0x0 ... fixed_offset + 0xf:
+#else
+	case fixed_offset + 0x0:
+	case fixed_offset + 0x1:
+	case fixed_offset + 0x2:
+	case fixed_offset + 0x3:
+	case fixed_offset + 0x4:
+	case fixed_offset + 0x5:
+	case fixed_offset + 0x6:
+	case fixed_offset + 0x7:
+	case fixed_offset + 0x8:
+	case fixed_offset + 0x9:
+	case fixed_offset + 0xa:
+	case fixed_offset + 0xb:
+	case fixed_offset + 0xc:
+	case fixed_offset + 0xd:
+	case fixed_offset + 0xe:
+	case fixed_offset + 0xf:
+#endif
+		++*off;
+		size = ((unsigned int)*p) & 0x0f;
+		break;
+	default:
+		PyErr_SetString(PyExc_ValueError, "Unexpected type header on stream");
+		return -1;
+    }
+	msgpack_unpack_callback(_uint32)(&ctx->user, size, &ctx->stack[0].obj);
+	return 1;
+}
+
+#undef SWITCH_RANGE_BEGIN
+#undef SWITCH_RANGE
+#undef SWITCH_RANGE_DEFAULT
+#undef SWITCH_RANGE_END
+
 static const execute_fn template_construct = &template_execute<true>;
 static const execute_fn template_skip = &template_execute<false>;
+static const execute_fn read_array_header = &template_container_header<0x90, 0xdc>;
+static const execute_fn read_map_header = &template_container_header<0x80, 0xde>;
 
 #undef msgpack_unpack_func
 #undef msgpack_unpack_callback
diff --git a/test/test_read_size.py b/test/test_read_size.py
new file mode 100644
index 0000000..714f963
--- /dev/null
+++ b/test/test_read_size.py
@@ -0,0 +1,66 @@
+"""Test Unpacker's read_array_header and read_map_header methods"""
+from msgpack import packb, Unpacker
+UnexpectedTypeException = ValueError
+
+def test_read_array_header():
+    unpacker = Unpacker()
+    unpacker.feed(packb(['a', 'b', 'c']))
+    assert unpacker.read_array_header() == 3
+    assert unpacker.unpack() == 'a'
+    assert unpacker.unpack() == 'b'
+    assert unpacker.unpack() == 'c'
+    try:
+        unpacker.unpack()
+        assert 0, 'should raise exception'
+    except StopIteration:
+        assert 1, 'okay'
+
+
+def test_read_map_header():
+    unpacker = Unpacker()
+    unpacker.feed(packb({'a': 'A'}))
+    assert unpacker.read_map_header() == 1
+    assert unpacker.unpack() == 'a'
+    assert unpacker.unpack() == 'A'
+    try:
+        unpacker.unpack()
+        assert 0, 'should raise exception'
+    except StopIteration:
+        assert 1, 'okay'
+
+def test_incorrect_type_array():
+    unpacker = Unpacker()
+    unpacker.feed(packb(1))
+    try:
+        unpacker.read_array_header()
+        assert 0, 'should raise exception'
+    except UnexpectedTypeException:
+        assert 1, 'okay'
+
+def test_incorrect_type_map():
+    unpacker = Unpacker()
+    unpacker.feed(packb(1))
+    try:
+        unpacker.read_map_header()
+        assert 0, 'should raise exception'
+    except UnexpectedTypeException:
+        assert 1, 'okay'
+
+def test_correct_type_nested_array():
+    unpacker = Unpacker()
+    unpacker.feed(packb({'a': ['b', 'c', 'd']}))
+    try:
+        unpacker.read_array_header()
+        assert 0, 'should raise exception'
+    except UnexpectedTypeException:
+        assert 1, 'okay'
+
+def test_incorrect_type_nested_map():
+    unpacker = Unpacker()
+    unpacker.feed(packb([{'a': 'b'}]))
+    try:
+        unpacker.read_map_header()
+        assert 0, 'should raise exception'
+    except UnexpectedTypeException:
+        assert 1, 'okay'
+