mirror of
https://github.com/msgpack/msgpack-c.git
synced 2025-03-19 21:18:23 +01:00
Add list_hook option to unpacker.
This commit is contained in:
parent
063d51c662
commit
e1711ffcf2
@ -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 = <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)
|
||||
|
@ -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)
|
||||
|
@ -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);*/
|
||||
|
@ -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()
|
||||
|
Loading…
x
Reference in New Issue
Block a user