Add list_hook option to unpacker.

methane · methane · commit 3d8978417a17 · 2010-10-26T02:09:52.000+09:00
diff --git a/msgpack/_msgpack.pyx b/msgpack/_msgpack.pyx
@@ -52,6 +52,7 @@ cdef extern from "pack.h":
     int msgpack_pack_raw(msgpack_packer* pk, size_t l)
     int msgpack_pack_raw_body(msgpack_packer* pk, char* body, size_t l)
 
+cdef int DEFAULT_RECURSE_LIMIT=511
 
 cdef class Packer(object):
     """MessagePack Packer
@@ -80,7 +81,8 @@ cdef class Packer(object):
     def __dealloc__(self):
         free(self.pk.buf);
 
-    cdef int _pack(self, object o, int nest_limit=511, default=None) except -1:
+    cdef int _pack(self, object o, int nest_limit=DEFAULT_RECURSE_LIMIT,
+                   default=None) except -1:
         cdef long long llval
         cdef unsigned long long ullval
         cdef long longval
@@ -148,7 +150,7 @@ cdef class Packer(object):
 
     def pack(self, object obj):
         cdef int ret
-        ret = self._pack(obj, self.default)
+        ret = self._pack(obj, DEFAULT_RECURSE_LIMIT, self.default)
         if ret:
             raise TypeError
         buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
@@ -172,6 +174,7 @@ cdef extern from "unpack.h":
     ctypedef struct msgpack_user:
         int use_list
         PyObject* object_hook
+        PyObject* list_hook
 
     ctypedef struct template_context:
         msgpack_user user
@@ -186,19 +189,23 @@ cdef extern from "unpack.h":
     object template_data(template_context* ctx)
 
 
-def unpackb(bytes packed_bytes, object object_hook=None):
+def unpackb(bytes packed_bytes, object object_hook=None, object list_hook=None):
     """Unpack packed_bytes to object. Returns an unpacked object."""
     cdef const_char_ptr p = packed_bytes
     cdef template_context ctx
     cdef size_t off = 0
     cdef int ret
     template_init(&ctx)
     ctx.user.use_list = 0
-    ctx.user.object_hook = NULL
+    ctx.user.object_hook = ctx.user.list_hook = NULL
     if object_hook is not None:
         if not PyCallable_Check(object_hook):
             raise TypeError("object_hook must be a callable.")
         ctx.user.object_hook = <PyObject*>object_hook
+    if list_hook is not None:
+        if not PyCallable_Check(list_hook):
+            raise TypeError("list_hook must be a callable.")
+        ctx.user.list_hook = <PyObject*>list_hook
     ret = template_execute(&ctx, p, len(packed_bytes), &off)
     if ret == 1:
         return template_data(&ctx)
@@ -207,10 +214,10 @@ def unpackb(bytes packed_bytes, object object_hook=None):
 
 unpacks = unpackb
 
-def unpack(object stream, object object_hook=None):
+def unpack(object stream, object object_hook=None, object list_hook=None):
     """unpack an object from stream."""
     packed = stream.read()
-    return unpackb(packed, object_hook=object_hook)
+    return unpackb(packed, object_hook=object_hook, list_hook=list_hook)
 
 cdef class UnpackIterator(object):
     cdef object unpacker
@@ -265,7 +272,7 @@ cdef class Unpacker(object):
             free(self.buf);
 
     def __init__(self, file_like=None, int read_size=0, bint use_list=0,
-                 object object_hook=None):
+                 object object_hook=None, object list_hook=None):
         if read_size == 0:
             read_size = 1024*1024
         self.use_list = use_list
@@ -278,11 +285,15 @@ cdef class Unpacker(object):
         self.buf_tail = 0
         template_init(&self.ctx)
         self.ctx.user.use_list = use_list
-        self.ctx.user.object_hook = <PyObject*>NULL
+        self.ctx.user.object_hook = self.ctx.user.list_hook = <PyObject*>NULL
         if object_hook is not None:
             if not PyCallable_Check(object_hook):
                 raise TypeError("object_hook must be a callable.")
             self.ctx.user.object_hook = <PyObject*>object_hook
+        if list_hook is not None:
+            if not PyCallable_Check(list_hook):
+                raise TypeError("object_hook must be a callable.")
+            self.ctx.user.list_hook = <PyObject*>list_hook
 
     def feed(self, bytes next_bytes):
         self.waiting_bytes.append(next_bytes)
diff --git a/msgpack/unpack.h b/msgpack/unpack.h
@@ -22,6 +22,7 @@
 typedef struct unpack_user {
     int use_list;
     PyObject *object_hook;
+    PyObject *list_hook;
 } unpack_user;
 
 
@@ -154,6 +155,16 @@ static inline int template_callback_array_item(unpack_user* u, unsigned int curr
     return 0;
 }
 
+static inline int template_callback_array_end(unpack_user* u, msgpack_unpack_object* c)
+{
+    if (u->list_hook) {
+        PyObject *arglist = Py_BuildValue("(O)", *c);
+        *c = PyEval_CallObject(u->list_hook, arglist);
+        Py_DECREF(arglist);
+    }
+    return 0;
+}
+
 static inline int template_callback_map(unpack_user* u, unsigned int n, msgpack_unpack_object* o)
 {
     PyObject *p = PyDict_New();
@@ -173,16 +184,14 @@ static inline int template_callback_map_item(unpack_user* u, msgpack_unpack_obje
     return -1;
 }
 
-//static inline int template_callback_map_end(unpack_user* u, msgpack_unpack_object* c)
-int template_callback_map_end(unpack_user* u, msgpack_unpack_object* c)
+static inline int template_callback_map_end(unpack_user* u, msgpack_unpack_object* c)
 {
     if (u->object_hook) {
         PyObject *arglist = Py_BuildValue("(O)", *c);
         *c = PyEval_CallObject(u->object_hook, arglist);
         Py_DECREF(arglist);
-        return 0;
     }
-    return -1;
+    return 0;
 }
 
 static inline int template_callback_raw(unpack_user* u, const char* b, const char* p, unsigned int l, msgpack_unpack_object* o)
diff --git a/msgpack/unpack_template.h b/msgpack/unpack_template.h
@@ -304,6 +304,7 @@ msgpack_unpack_func(int, _execute)(msgpack_unpack_struct(_context)* ctx, const c
 	case CT_ARRAY_ITEM:
 		if(msgpack_unpack_callback(_array_item)(user, c->curr, &c->obj, obj) < 0) { goto _failed; }
 		if(++c->curr == c->count) {
+			msgpack_unpack_callback(_array_end)(user, &c->obj);
 			obj = c->obj;
 			--top;
 			/*printf("stack pop %d\n", top);*/
@@ -317,7 +318,7 @@ msgpack_unpack_func(int, _execute)(msgpack_unpack_struct(_context)* ctx, const c
 	case CT_MAP_VALUE:
 		if(msgpack_unpack_callback(_map_item)(user, &c->obj, c->map_key, obj) < 0) { goto _failed; }
 		if(--c->count == 0) {
-            msgpack_unpack_callback(_map_end)(user, &c->obj);
+			msgpack_unpack_callback(_map_end)(user, &c->obj);
 			obj = c->obj;
 			--top;
 			/*printf("stack pop %d\n", top);*/
diff --git a/test/test_obj.py b/test/test_obj.py
@@ -31,7 +31,16 @@ def test_bad_hook():
     packed = packs([3, 1+2j], default=lambda o: o)
     unpacked = unpacks(packed)
 
+def _arr_to_str(arr):
+    return ''.join(str(c) for c in arr)
+
+def test_array_hook():
+    packed = packs([1,2,3])
+    unpacked = unpacks(packed, list_hook=_arr_to_str)
+    eq_(unpacked, '123')
+
 if __name__ == '__main__':
     test_decode_hook()
     test_encode_hook()
     test_bad_hook()
+    test_array_hook()
