From e1711ffcf203d43ee71cfb5eea8e65f3e50b145b Mon Sep 17 00:00:00 2001 From: INADA Naoki Date: Tue, 26 Oct 2010 02:09:52 +0900 Subject: [PATCH] Add list_hook option to unpacker. --- python/msgpack/_msgpack.pyx | 27 +++++++++++++++++++-------- python/msgpack/unpack.h | 17 +++++++++++++---- python/msgpack/unpack_template.h | 3 ++- python/test/test_obj.py | 9 +++++++++ 4 files changed, 43 insertions(+), 13 deletions(-) diff --git a/python/msgpack/_msgpack.pyx b/python/msgpack/_msgpack.pyx index 0abdd515..e9d6c7bc 100644 --- a/python/msgpack/_msgpack.pyx +++ b/python/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,7 +189,7 @@ 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 @@ -194,11 +197,15 @@ def unpackb(bytes packed_bytes, object object_hook=None): 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 = 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 = 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 = NULL + self.ctx.user.object_hook = self.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.") self.ctx.user.object_hook = 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 = list_hook def feed(self, bytes next_bytes): self.waiting_bytes.append(next_bytes) diff --git a/python/msgpack/unpack.h b/python/msgpack/unpack.h index 404ee5ad..453ec2b8 100644 --- a/python/msgpack/unpack.h +++ b/python/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/python/msgpack/unpack_template.h b/python/msgpack/unpack_template.h index 1fdedd70..7a2288f1 100644 --- a/python/msgpack/unpack_template.h +++ b/python/msgpack/unpack_template.h @@ -304,6 +304,7 @@ _push: 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 @@ _push: 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/python/test/test_obj.py b/python/test/test_obj.py index 28edacb8..bc857361 100644 --- a/python/test/test_obj.py +++ b/python/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()