Add list_hook option to unpacker.

This commit is contained in:
INADA Naoki 2010-10-26 02:09:52 +09:00
parent 063d51c662
commit e1711ffcf2
4 changed files with 43 additions and 13 deletions

View File

@ -52,6 +52,7 @@ cdef extern from "pack.h":
int msgpack_pack_raw(msgpack_packer* pk, size_t l) int msgpack_pack_raw(msgpack_packer* pk, size_t l)
int msgpack_pack_raw_body(msgpack_packer* pk, char* body, 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): cdef class Packer(object):
"""MessagePack Packer """MessagePack Packer
@ -80,7 +81,8 @@ cdef class Packer(object):
def __dealloc__(self): def __dealloc__(self):
free(self.pk.buf); 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 long long llval
cdef unsigned long long ullval cdef unsigned long long ullval
cdef long longval cdef long longval
@ -148,7 +150,7 @@ cdef class Packer(object):
def pack(self, object obj): def pack(self, object obj):
cdef int ret cdef int ret
ret = self._pack(obj, self.default) ret = self._pack(obj, DEFAULT_RECURSE_LIMIT, self.default)
if ret: if ret:
raise TypeError raise TypeError
buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length) buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
@ -172,6 +174,7 @@ cdef extern from "unpack.h":
ctypedef struct msgpack_user: ctypedef struct msgpack_user:
int use_list int use_list
PyObject* object_hook PyObject* object_hook
PyObject* list_hook
ctypedef struct template_context: ctypedef struct template_context:
msgpack_user user msgpack_user user
@ -186,7 +189,7 @@ cdef extern from "unpack.h":
object template_data(template_context* ctx) 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.""" """Unpack packed_bytes to object. Returns an unpacked object."""
cdef const_char_ptr p = packed_bytes cdef const_char_ptr p = packed_bytes
cdef template_context ctx cdef template_context ctx
@ -194,11 +197,15 @@ def unpackb(bytes packed_bytes, object object_hook=None):
cdef int ret cdef int ret
template_init(&ctx) template_init(&ctx)
ctx.user.use_list = 0 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 object_hook is not None:
if not PyCallable_Check(object_hook): if not PyCallable_Check(object_hook):
raise TypeError("object_hook must be a callable.") raise TypeError("object_hook must be a callable.")
ctx.user.object_hook = <PyObject*>object_hook 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) ret = template_execute(&ctx, p, len(packed_bytes), &off)
if ret == 1: if ret == 1:
return template_data(&ctx) return template_data(&ctx)
@ -207,10 +214,10 @@ def unpackb(bytes packed_bytes, object object_hook=None):
unpacks = unpackb 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.""" """unpack an object from stream."""
packed = stream.read() 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 class UnpackIterator(object):
cdef object unpacker cdef object unpacker
@ -265,7 +272,7 @@ cdef class Unpacker(object):
free(self.buf); free(self.buf);
def __init__(self, file_like=None, int read_size=0, bint use_list=0, 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: if read_size == 0:
read_size = 1024*1024 read_size = 1024*1024
self.use_list = use_list self.use_list = use_list
@ -278,11 +285,15 @@ cdef class Unpacker(object):
self.buf_tail = 0 self.buf_tail = 0
template_init(&self.ctx) template_init(&self.ctx)
self.ctx.user.use_list = use_list 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 object_hook is not None:
if not PyCallable_Check(object_hook): if not PyCallable_Check(object_hook):
raise TypeError("object_hook must be a callable.") raise TypeError("object_hook must be a callable.")
self.ctx.user.object_hook = <PyObject*>object_hook 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): def feed(self, bytes next_bytes):
self.waiting_bytes.append(next_bytes) self.waiting_bytes.append(next_bytes)

View File

@ -22,6 +22,7 @@
typedef struct unpack_user { typedef struct unpack_user {
int use_list; int use_list;
PyObject *object_hook; PyObject *object_hook;
PyObject *list_hook;
} unpack_user; } unpack_user;
@ -154,6 +155,16 @@ static inline int template_callback_array_item(unpack_user* u, unsigned int curr
return 0; 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) static inline int template_callback_map(unpack_user* u, unsigned int n, msgpack_unpack_object* o)
{ {
PyObject *p = PyDict_New(); PyObject *p = PyDict_New();
@ -173,16 +184,14 @@ static inline int template_callback_map_item(unpack_user* u, msgpack_unpack_obje
return -1; return -1;
} }
//static inline 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)
int template_callback_map_end(unpack_user* u, msgpack_unpack_object* c)
{ {
if (u->object_hook) { if (u->object_hook) {
PyObject *arglist = Py_BuildValue("(O)", *c); PyObject *arglist = Py_BuildValue("(O)", *c);
*c = PyEval_CallObject(u->object_hook, arglist); *c = PyEval_CallObject(u->object_hook, arglist);
Py_DECREF(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) static inline int template_callback_raw(unpack_user* u, const char* b, const char* p, unsigned int l, msgpack_unpack_object* o)

View File

@ -304,6 +304,7 @@ _push:
case CT_ARRAY_ITEM: case CT_ARRAY_ITEM:
if(msgpack_unpack_callback(_array_item)(user, c->curr, &c->obj, obj) < 0) { goto _failed; } if(msgpack_unpack_callback(_array_item)(user, c->curr, &c->obj, obj) < 0) { goto _failed; }
if(++c->curr == c->count) { if(++c->curr == c->count) {
msgpack_unpack_callback(_array_end)(user, &c->obj);
obj = c->obj; obj = c->obj;
--top; --top;
/*printf("stack pop %d\n", top);*/ /*printf("stack pop %d\n", top);*/
@ -317,7 +318,7 @@ _push:
case CT_MAP_VALUE: case CT_MAP_VALUE:
if(msgpack_unpack_callback(_map_item)(user, &c->obj, c->map_key, obj) < 0) { goto _failed; } if(msgpack_unpack_callback(_map_item)(user, &c->obj, c->map_key, obj) < 0) { goto _failed; }
if(--c->count == 0) { if(--c->count == 0) {
msgpack_unpack_callback(_map_end)(user, &c->obj); msgpack_unpack_callback(_map_end)(user, &c->obj);
obj = c->obj; obj = c->obj;
--top; --top;
/*printf("stack pop %d\n", top);*/ /*printf("stack pop %d\n", top);*/

View File

@ -31,7 +31,16 @@ def test_bad_hook():
packed = packs([3, 1+2j], default=lambda o: o) packed = packs([3, 1+2j], default=lambda o: o)
unpacked = unpacks(packed) 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__': if __name__ == '__main__':
test_decode_hook() test_decode_hook()
test_encode_hook() test_encode_hook()
test_bad_hook() test_bad_hook()
test_array_hook()