Merge commit from fork

Author: methane

msgpack/_unpacker.pyx

@@ -40,20 +40,22 @@ cdef extern from "unpack.h":
         Py_ssize_t max_map_len
         Py_ssize_t max_ext_len
 
+    ctypedef struct unpack_context
+    ctypedef int (*execute_fn)(unpack_context* ctx, const char* data,
+                               Py_ssize_t len, Py_ssize_t* off) except -1
+
     ctypedef struct unpack_context:
         msgpack_user user
-        PyObject* obj
-        Py_ssize_t count
+        execute_fn execute
 
-    ctypedef int (*execute_fn)(unpack_context* ctx, const char* data,
-                               Py_ssize_t len, Py_ssize_t* off) except -1
     execute_fn unpack_construct
     execute_fn unpack_skip
     execute_fn read_array_header
     execute_fn read_map_header
 
     void unpack_init(unpack_context* ctx)
-    object unpack_data(unpack_context* ctx)
+    # Transfers the result's owned reference to Cython and clears its slot.
+    object unpack_take_result(unpack_context* ctx)
     void unpack_clear(unpack_context* ctx)
 
 cdef inline init_ctx(unpack_context *ctx,
@@ -191,16 +193,16 @@ def unpackb(object packed, *, object object_hook=None, object list_hook=None,
                  max_str_len, max_bin_len, max_array_len, max_map_len, max_ext_len)
         ret = unpack_construct(&ctx, buf, buf_len, &off)
         if ret == 1:
-            obj = unpack_data(&ctx)
+            obj = unpack_take_result(&ctx)
             if off < buf_len:
                 # buf may point into a temporary contiguous copy owned by view,
                 # so the extra data must be copied out before releasing view.
                 raise ExtraData(obj, PyBytes_FromStringAndSize(buf+off, buf_len-off))
             return obj
     finally:
+        unpack_clear(&ctx)
         PyBuffer_Release(&view);
 
-    unpack_clear(&ctx)
     if ret == 0:
         raise ValueError("Unpack failed: incomplete input")
     elif ret == -2:
@@ -214,6 +216,11 @@ def unpackb(object packed, *, object object_hook=None, object list_hook=None,
 cdef class Unpacker:
     """Streaming unpacker.
 
+    If an operation needs more data, resume it with the same method after
+    feeding more bytes. ``unpack()`` and iteration may be used interchangeably.
+    Switching between unpacking, skipping, header readers, or ``read_bytes()``
+    while an object is incomplete raises ``ValueError`` in the C extension.
+
     Arguments:
 
     :param file_like:
@@ -339,6 +346,9 @@ cdef class Unpacker:
                  Py_ssize_t max_ext_len=-1):
         cdef const char *cerr=NULL
 
+        if self._unpacking:
+            raise RuntimeError("Unpacker.__init__() cannot be called while unpacking is in progress")
+
         unpack_clear(&self.ctx)
         unpack_init(&self.ctx)
         if self.buf != NULL:
@@ -473,18 +483,29 @@ cdef class Unpacker:
         cdef object obj
         cdef Py_ssize_t prev_head
 
+        if self._unpacking:
+            raise RuntimeError("Unpacker cannot be called recursively while unpacking is in progress")
+        if self.ctx.execute != NULL and self.ctx.execute != execute:
+            raise ValueError("Cannot switch unpacking methods while an object is incomplete")
+
         self._unpacking = True
         try:
             while 1:
                 prev_head = self.buf_head
                 if prev_head < self.buf_tail:
-                    ret = execute(&self.ctx, self.buf, self.buf_tail, &self.buf_head)
-                    self.stream_offset += self.buf_head - prev_head
+                    self.ctx.execute = execute
+                    try:
+                        ret = execute(&self.ctx, self.buf, self.buf_tail, &self.buf_head)
+                    except:
+                        unpack_clear(&self.ctx)
+                        raise
+                    finally:
+                        self.stream_offset += self.buf_head - prev_head
                 else:
                     ret = 0
 
                 if ret == 1:
-                    obj = unpack_data(&self.ctx)
+                    obj = unpack_take_result(&self.ctx)
                     unpack_init(&self.ctx)
                     return obj
                 if ret == 0:
@@ -510,6 +531,10 @@ cdef class Unpacker:
     def read_bytes(self, Py_ssize_t nbytes):
         """Read a specified number of raw bytes from the stream"""
         cdef Py_ssize_t nread
+        if self._unpacking:
+            raise RuntimeError("Unpacker.read_bytes() cannot be called while unpacking is in progress")
+        if self.ctx.execute != NULL:
+            raise ValueError("Cannot switch unpacking methods while an object is incomplete")
         nread = min(self.buf_tail - self.buf_head, nbytes)
         ret = PyBytes_FromStringAndSize(self.buf + self.buf_head, nread)
         self.buf_head += nread

msgpack/unpack.h

@@ -35,11 +35,13 @@ typedef struct unpack_user {
     Py_ssize_t max_str_len, max_bin_len, max_array_len, max_map_len, max_ext_len;
 } unpack_user;
 
-typedef PyObject* msgpack_unpack_object;
+typedef PyObject* msgpack_unpack_object; //TODO: remove typedef and use PyObject* directly
 struct unpack_context;
 typedef struct unpack_context unpack_context;
 typedef int (*execute_fn)(unpack_context *ctx, const char* data, Py_ssize_t len, Py_ssize_t* off);
 
+/* Constructors receive a NULL output slot and store an owned reference on
+ * success. On failure the slot remains NULL. */
 static inline int unpack_callback_uint16(unpack_user* u, uint16_t d, msgpack_unpack_object* o)
 {
     PyObject *p = PyLong_FromLong((long)d);
@@ -137,6 +139,7 @@ static inline int unpack_callback_array(unpack_user* u, unsigned int n, msgpack_
     return 0;
 }
 
+/* Consumes o on both success and failure; borrows the container. */
 static inline int unpack_callback_array_item(unpack_user* u, unsigned int current, msgpack_unpack_object* c, msgpack_unpack_object o)
 {
     if (u->use_list)
@@ -146,14 +149,16 @@ static inline int unpack_callback_array_item(unpack_user* u, unsigned int curren
     return 0;
 }
 
+/* Replaces the owned container reference with the hook result, or NULL on
+ * failure. The original reference is consumed in either case. */
 static inline int unpack_callback_array_end(unpack_user* u, msgpack_unpack_object* c)
 {
     if (u->list_hook) {
         PyObject *new_c = PyObject_CallFunctionObjArgs(u->list_hook, *c, NULL);
-        if (!new_c)
-            return -1;
         Py_DECREF(*c);
         *c = new_c;
+        if (!new_c)
+            return -1;
     }
     return 0;
 }
@@ -177,41 +182,42 @@ static inline int unpack_callback_map(unpack_user* u, unsigned int n, msgpack_un
     return 0;
 }
 
+/* Consumes k and v on both success and failure; borrows the container. */
 static inline int unpack_callback_map_item(unpack_user* u, unsigned int current, msgpack_unpack_object* c, msgpack_unpack_object k, msgpack_unpack_object v)
 {
     if (u->strict_map_key && !PyUnicode_CheckExact(k) && !PyBytes_CheckExact(k)) {
         PyErr_Format(PyExc_ValueError, "%.100s is not allowed for map key when strict_map_key=True", Py_TYPE(k)->tp_name);
+        Py_DECREF(k);
+        Py_DECREF(v);
         return -1;
     }
     if (PyUnicode_CheckExact(k)) {
         PyUnicode_InternInPlace(&k);
     }
     if (u->has_pairs_hook) {
         msgpack_unpack_object item = PyTuple_Pack(2, k, v);
-        if (!item)
-            return -1;
         Py_DECREF(k);
         Py_DECREF(v);
+        if (!item)
+            return -1;
         PyList_SET_ITEM(*c, current, item);
         return 0;
     }
-    else if (PyDict_SetItem(*c, k, v) == 0) {
-        Py_DECREF(k);
-        Py_DECREF(v);
-        return 0;
-    }
-    return -1;
+    int ret = PyDict_SetItem(*c, k, v);
+    Py_DECREF(k);
+    Py_DECREF(v);
+    return ret;
 }
 
+/* Same ownership contract as unpack_callback_array_end. */
 static inline int unpack_callback_map_end(unpack_user* u, msgpack_unpack_object* c)
 {
     if (u->object_hook) {
         PyObject *new_c = PyObject_CallFunctionObjArgs(u->object_hook, *c, NULL);
-        if (!new_c)
-            return -1;
-
         Py_DECREF(*c);
         *c = new_c;
+        if (!new_c)
+            return -1;
     }
     return 0;
 }

msgpack/unpack_container_header.h

@@ -46,9 +46,8 @@ static inline int unpack_container_header(unpack_context* ctx, const char* data,
         return -1;
     }
 
-    if (unpack_callback_uint32(&ctx->user, size, &ctx->stack[0].obj) < 0)
+    if (unpack_callback_uint32(&ctx->user, size, &ctx->result) < 0)
         return -1;
 
     return 1;
 }
-