8000 Support object_pairs_hook · peter80/msgpack-python@e381032 · GitHub
[go: up one dir, main page]
Skip to content

Commit e381032

Browse files
methanemethane
committed
Support object_pairs_hook
Merge remote-tracking branch 'jnothman/object_pairs_hook' into 0.2-maint Conflicts: msgpack/_msgpack.pyx test/test_pack.py test/test_sequnpack.py
2 parents 927d291 + 7794251 commit e381032

File tree

7 files changed

+143
-90
lines changed

7 files changed

+143
-90
lines changed

‎ChangeLog.rst

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,12 @@
1+
0.3.0
2+
=====
3+
:release date: in development
4+
5+
Changes
6+
-------
7+
* Add ``.skip()`` method to ``Unpacker`` (thanks to jnothman)
8+
9+
110
0.2.3
211
=======
312
:release date: in development

‎msgpack/_msgpack.pyx

Lines changed: 71 additions & 67 deletions
467
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,7 @@ cdef extern from "unpack.h":
201201
ctypedef struct msgpack_user:
202202
bint use_list
203203
PyObject* object_hook
204+
bint has_pairs_hook # call object_hook with k-v pairs
204205
PyObject* list_hook
205206
char *encoding
206207
char *unicode_errors
@@ -213,13 +214,54 @@ cdef extern from "unpack.h":
213214
PyObject* key
214215

215216
int template_execute(template_context* ctx, const_char_ptr data,
216-
size_t len, size_t* off) except -1
217+
size_t len, size_t* off, bint construct) except -1
217218
void template_init(template_context* ctx)
218219
object template_data(template_context* ctx)
219220

221+
cdef inline init_ctx(template_context *ctx, object object_hook, object object_pairs_hook, object list_hook, bint use_list, encoding, unicode_errors):
222+
template_init(ctx)
223+
ctx.user.use_list = use_list
224+
ctx.user.object_hook = ctx.user.list_hook = <PyObject*>NULL
225+
226+
if object_hook is not None and object_pairs_hook is not None:
227+
raise ValueError("object_pairs_hook and object_hook are mutually exclusive.")
228+
229+
if object_hook is not None:
230+
if not PyCallable_Check(object_hook):
231+
raise TypeError("object_hook must be a callable.")
232+
ctx.user.object_hook = <PyObject*>object_hook
233+
234+
if object_pairs_hook is None:
235+
ctx.user.has_pairs_hook = False
236+
else:
237+
if not PyCallable_Check(object_pairs_hook):
238+
raise TypeError("object_pairs_hook must be a callable.")
239+
ctx.user.object_hook = <PyObject*>object_pairs_hook
240+
ctx.user.has_pairs_hook = True
241+
242+
if list_hook is not None:
243+
if not PyCallable_Check(list_hook):
244+
raise TypeError("list_hook must be a callable.")
245+
ctx.user.list_hook = <PyObject*>list_hook
246+
247+
if encoding is None:
248+
ctx.user.encoding = NULL
249+
ctx.user.unicode_errors = NULL
250+
else:
251+
if isinstance(encoding, unicode):
252+
_bencoding = encoding.encode('ascii')
253+
else:
254+
_bencoding = encoding
255+
ctx.user.encoding = PyBytes_AsString(_bencoding)
256+
if isinstance(unicode_errors, unicode):
257+
_berrors = unicode_errors.encode('ascii')
258+
else:
259+
_berrors = unicode_errors
260+
ctx.user.unicode_errors = PyBytes_AsString(_berrors)
220261

221262
def unpackb(object packed, object object_hook=None, object list_hook=None,
222263
use_list=None, encoding=None, unicode_errors="strict",
264+
object_pairs_hook=None,
223265
):
224266
"""Unpack packed_bytes to object. Returns an unpacked object.
225267
@@ -234,39 +276,11 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
234276

235277
PyObject_AsReadBuffer(packed, <const_void_ptr*>&buf, &buf_len)
236278

237-
if encoding is None:
238-
enc = NULL
239-
err = NULL
240-
else:
241-
if isinstance(encoding, unicode):
242-
bencoding = encoding.encode('ascii')
243-
else:
244-
bencoding = encoding
245-
if isinstance(unicode_errors, unicode):
246-
berrors = unicode_errors.encode('ascii')
247-
else:
248-
berrors = unicode_errors
249-
enc = PyBytes_AsString(bencoding)
250-
err = PyBytes_AsString(berrors)
251-
252-
template_init(&ctx)
253279
if use_list is None:
254280
warnings.warn("Set use_list explicitly.", category=DeprecationWarning, stacklevel=1)
255-
ctx.user.use_list = 0
256-
else:
257-
ctx.user.use_list = use_list
258-
ctx.user.object_hook = ctx.user.list_hook = NULL
259-
ctx.user.encoding = <const_char_ptr>enc
260-
ctx.user.unicode_errors = <const_char_ptr>err
261-
if object_hook is not None:
262-
if not PyCallable_Check(object_hook):
263-
raise TypeError("object_hook must be a callable.")
264-
ctx.user.object_hook = <PyObject*>object_hook
265-
if list_hook is not None:
266-
if not PyCallable_Check(list_hook):
267-
raise TypeError("list_hook must be a callable.")
268-
ctx.user.list_hook = <PyObject*>list_hook
269-
ret = template_execute(&ctx, buf, buf_len, &off)
281+
use_list = 0
282+
init_ctx(&ctx, object_hook, object_pairs_hook, list_hook, use_list, encoding, unicode_errors)
283+
ret = template_execute(&ctx, buf, buf_len, &off, 1)
270284
if ret == 1:
271285
obj = template_data(&ctx)
272286
if off < buf_len:
@@ -278,6 +292,7 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
278292

279293
def unpack(object stream, object object_hook=None, object list_hook=None,
280294
use_list=None, encoding=None, unicode_errors="strict",
295+
object_pairs_hook=None,
281296
):
282297
"""Unpack an object from `stream`.
283298
@@ -287,7 +302,7 @@ def unpack(object stream, object object_hook=None, object list_hook=None,
287302
warnings.warn("Set use_list explicitly.", category=DeprecationWarning, stacklevel=1)
288303
use_list = 0
289304
return unpackb(stream.read(), use_list=use_list,
290-
object_hook=object_hook, list_hook=list_hook,
305+
object_hook=object_hook, object_pairs_hook=object_pairs_hook, list_hook=list_hook,
291306
encoding=encoding, unicode_errors=unicode_errors,
292307
)
293308

@@ -307,7 +322,10 @@ cdef class Unpacker(object):
307322
Otherwise, it is deserialized to Python tuple.
308323
309324
`object_hook` is same to simplejson. If it is not None, it should be callable
310-
and Unpacker calls it when deserializing key-value.
325+
and Unpacker calls it with a dict argument after deserializing a map.
326+
327+
`object_pairs_hook` is same to simplejson. If it is not None, it should be callable
328+
and Unpacker calls it with a list of key-value pairs after deserializing a map.
311329
312330
`encoding` is encoding used for decoding msgpack bytes. If it is None (default),
313331
msgpack bytes is deserialized to Python bytes.
@@ -357,9 +375,8 @@ cdef class Unpacker(object):
357375
self.buf = NULL
358376

359377
def __init__(self, file_like=None, Py_ssize_t read_size=0, use_list=None,
360-
object object_hook=None, object list_hook=None,
378+
object object_hook=None, object object_pairs_hook=None, object list_hook=None,
361379
encoding=None, unicode_errors='strict', int max_buffer_size=0,
362-
object object_pairs_hook=None,
363380
):
364381
if use_list is None:
365382
warnings.warn("Set use_list explicitly.", category=DeprecationWarning, stacklevel=1)
@@ -384,31 +401,7 @@ cdef class Unpacker(object):
384401
self.buf_size = read_size
385402
self.buf_head = 0
386403
self.buf_tail = 0
387-
template_init(&self.ctx)
388-
self.ctx.user.use_list = use_list
389-
self.ctx.user.object_hook = self.ctx.user.list_hook = <PyObject*>NULL
390-
if object_hook is not None:
391-
if not PyCallable_Check(object_hook):
392-
raise TypeError("object_hook must be a callable.")
393-
self.ctx.user.object_hook = <PyObject*>object_hook
394-
if list_hook is not None:
395-
if not PyCallable_Check(list_hook):
396-
raise TypeError("list_hook must be a callable.")
397-
self.ctx.user.list_hook = <PyObject*>list_hook
398-
if encoding is None:
399-
self.ctx.user.encoding = NULL
400-
self.ctx.user.unicode_errors = NULL
401-
else:
402-
if isinstance(encoding, unicode):
403-
self._bencoding = encoding.encode('ascii')
404-
else:
405-
self._bencoding = encoding
406-
self.ctx.user.encoding = PyBytes_AsString(self._bencoding)
407-
if isinstance(unicode_errors, unicode):
408-
self._berrors = unicode_errors.encode('ascii')
409-
else:
410-
self._berrors = unicode_errors
411-
self.ctx.user.unicode_errors = PyBytes_AsString(self._berrors)
404+
init_ctx(&self.ctx, object_hook, object_pairs_hook, list_hook, use_list, encoding, unicode_errors)
412405

413406
def feed(self, object next_bytes):
414407
cdef char* buf
@@ -469,15 +462,18 @@ cdef class Unpacker(object):
469462
else:
470463
self.file_like = None
471464

472-
cpdef unpack(self):
473-
"""unpack one object"""
465+
cdef object _unpack(self, bint construct):
474466
cdef int ret
+
cdef object obj
475468
while 1:
476-
ret = template_execute(&self.ctx, self.buf, self.buf_tail, &self.buf_head)
469+
ret = template_execute(&self.ctx, self.buf, self.buf_tail, &self.buf_head, construct)
477470
if ret == 1:
478-
o = template_data(&self.ctx)
471+
if construct:
472+
obj = template_data(&self.ctx)
473+
else:
474+
obj = None
479475
template_init(&self.ctx)
480-
return o
476+
return obj
481477
elif ret == 0:
482478
if self.file_like is not None:
483479
self.read_from_file()
@@ -486,11 +482,19 @@ cdef class Unpacker(object):
486482
else:
487483
raise ValueError("Unpack failed: error = %d" % (ret,))
488484

485+
def unpack(self):
486+
"""unpack one object"""
487+
return self._unpack(1)
488+
489+
def skip(self):
490+
"""read and ignore one object, returning None"""
491+
return self._unpack(0)
492+
489493
def __iter__(self):
490494
return self
491495

492496
def __next__(self):
493-
return self.unpack()
497+
return self._unpack(1)
494498

495499
# for debug.
496500
#def _buf(self):

‎msgpack/unpack.h

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
typedef struct unpack_user {
2323
int use_list;
2424
PyObject *object_hook;
25+
bool has_pairs_hook;
2526
PyObject *list_hook;
2627
const char *encoding;
2728
const char *unicode_errors;
@@ -160,9 +161,7 @@ static inline int template_callback_array_item(unpack_user* u, unsigned int curr
160161
static inline int template_callback_array_end(unpack_user* u, msgpack_unpack_object* c)
161162
{
162163
if (u->list_hook) {
163-
PyObject *arglist = Py_BuildValue("(O)", *c);
164-
PyObject *new_c = PyEval_CallObject(u->list_hook, arglist);
165-
Py_DECREF(arglist);
164+
PyObject *new_c = PyEval_CallFunction(u->list_hook, "(O)", *c);
166165
Py_DECREF(*c);
167166
*c = new_c;
168167
}
@@ -171,16 +170,31 @@ static inline int template_callback_array_end(unpack_user* u, msgpack_unpack_obj
171170

172171
static inline int template_callback_map(unpack_user* u, unsigned int n, msgpack_unpack_object* o)
173172
{
174-
PyObject *p = PyDict_New();
173+
PyObject *p;
174+
if (u->has_pairs_hook) {
175+
p = PyList_New(n); // Or use tuple?
176+
}
177+
else {
178+
p = PyDict_New();
179+
}
175180
if (!p)
176181
return -1;
177182
*o = p;
178183
return 0;
179184
}
180185

181-
static inline int template_callback_map_item(unpack_user* u, msgpack_unpack_object* c, msgpack_unpack_object k, msgpack_unpack_object v)
186+
static inline int template_callback_map_item(unpack_user* u, unsigned int current, msgpack_unpack_object* c, msgpack_unpack_object k, msgpack_unpack_object v)
182187
{
183-
if (PyDict_SetItem(*c, k, v) == 0) {
188+
if (u->has_pairs_hook) {
189+
msgpack_unpack_object item = PyTuple_Pack(2, k, v);
190+
if (!item)
191+
return -1;
192+
Py_DECREF(k);
193+
Py_DECREF(v);
194+
PyList_SET_ITEM(*c, current, item);
195+
return 0;
196+
}
197+
else if (PyDict_SetItem(*c, k, v) == 0) {
184198
Py_DECREF(k);
185199
Py_DECREF(v);
186200
return 0;
@@ -191,9 +205,7 @@ static inline int template_callback_map_item(unpack_user* u, msgpack_unpack_obje
191205
static inline int template_callback_map_end(unpack_user* u, msgpack_unpack_object* c)
192206
{
193207
if (u->object_hook) {
194-
PyObject *arglist = Py_BuildValue("(O)", *c);
195-
PyObject *new_c = PyEval_CallObject(u->object_hook, arglist);
196-
Py_DECREF(arglist);
208+
PyObject *new_c = PyEval_CallFunction(u->object_hook, "(O)", *c);
197209
Py_DECREF(*c);
198210
*c = new_c;
199211
}

‎msgpack/unpack_template.h

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ msgpack_unpack_func(msgpack_unpack_object, _data)(msgpack_unpack_struct(_context
9595
}
9696

9797

98-
msgpack_unpack_func(int, _execute)(msgpack_unpack_struct(_context)* ctx, const char* data, size_t len, size_t* off)
98+
msgpack_unpack_func(int, _execute)(msgpack_unpack_struct(_context)* ctx, const char* data, size_t len, size_t* off, int construct)
9999
{
100100
assert(len >= *off);
101101

@@ -117,14 +117,17 @@ msgpack_unpack_func(int, _execute)(msgpack_unpack_struct(_context)* ctx, const c
117117

118118
int ret;
119119

120+
#define construct_cb(name) \
121+
construct && msgpack_unpack_callback(name)
122+
120123
#define push_simple_value(func) \
121-
if(msgpack_unpack_callback(func)(user, &obj) < 0) { goto _failed; } \
124+
if(construct_cb(func)(user, &obj) < 0) { goto _failed; } \
122125
goto _push
123126
#define push_fixed_value(func, arg) \
124-
if(msgpack_unpack_callback(func)(user, arg, &obj) < 0) { goto _failed; } \
127+
if(construct_cb(func)(user, arg, &obj) < 0) { goto _failed; } \
125128
goto _push
126129
#define push_variable_value(func, base, pos, len) \
127-
if(msgpack_unpack_callback(func)(user, \
130+
if(construct_cb(func)(user, \
128131
(const char*)base, (const char*)pos, len, &obj) < 0) { goto _failed; } \
129132
goto _push
130133

@@ -140,9 +143,9 @@ msgpack_unpack_func(int, _execute)(msgpack_unpack_struct(_context)* ctx, const c
140143

141144
#define start_container(func, count_, ct_) \
142145
if(top >= MSGPACK_EMBED_STACK_SIZE) { goto _failed; } /* FIXME */ \
143-
if(msgpack_unpack_callback(func)(user, count_, &stack[top].obj) < 0) { goto _failed; } \
146+
if(construct_cb(func)(user, count_, &stack[top].obj) < 0) { goto _failed; } \
144147
if((count_) == 0) { obj = stack[top].obj; \
145-
msgpack_unpack_callback(func##_end)(user, &obj); \
148+
construct_cb(func##_end)(user, &obj); \
146149
goto _push; } \
147150
stack[top].ct = ct_; \
148151
stack[top].size = count_; \
@@ -340,10 +343,10 @@ msgpack_unpack_func(int, _execute)(msgpack_unpack_struct(_context)* ctx, const c
340343
c = &stack[top-1];
341344
switch(c->ct) {
342345
case CT_ARRAY_ITEM:
343-
if(msgpack_unpack_callback(_array_item)(user, c->count, &c->obj, obj) < 0) { goto _failed; }
346+
if(construct_cb(_array_item)(user, c->count, &c->obj, obj) < 0) { goto _failed; }
344347
if(++c->count == c->size) {
345348
obj = c->obj;
346-
msgpack_unpack_callback(_array_end)(user, &obj);
349+
construct_cb(_array_end)(user, &obj);
347350
--top;
348351
/*printf("stack pop %d\n", top);*/
349352
goto _push;
@@ -354,10 +357,10 @@ msgpack_unpack_func(int, _execute)(msgpack_unpack_struct(_context)* ctx, const c
354357
c->ct = CT_MAP_VALUE;
355358
goto _header_again;
356359
case CT_MAP_VALUE:
357-
if(msgpack_unpack_callback(_map_item)(user, &c->obj, c->map_key, obj) < 0) { goto _failed; }
360+
if(construct_cb(_map_item)(user, c->count, &c->obj, c->map_key, obj) < 0) { goto _failed; }
358361
if(++c->count == c->size) {
359362
obj = c->obj;
360-
msgpack_unpack_callback(_map_end)(user, &obj);
363+
construct_cb(_map_end)(user, &obj);
361364
--top;
362365
/*printf("stack pop %d\n", top);*/
363366
goto _push;
@@ -399,6 +402,7 @@ msgpack_unpack_func(int, _execute)(msgpack_unpack_struct(_context)* ctx, const c
399402
*off = p - (const unsigned char*)data;
400403

401404
return ret;
405+
#undef construct_cb
402406
}
403407

404408

‎test/test_obj.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,16 @@ def test_decode_hook():
2626
unpacked = unpackb(packed, object_hook=_decode_complex, use_list=1)
2727
eq_(unpacked[1], 1+2j)
2828

29+
def test_decode_pairs_hook():
30+
packed = packb([3, {1: 2, 3: 4}])
31+
prod_sum = 1 * 2 + 3 * 4
32+
unpacked = unpackb(packed, object_pairs_hook=lambda l: sum(k * v for k, v in l))
33+
eq_(unpacked[1], prod_sum)
34+
35+
@raises(ValueError)
36+
def test_only_one_obj_hook():
37+
unpackb(b'', object_hook=lambda x: x, object_pairs_hook=lambda x: x)
38+
2939
@raises(ValueError)
3040
def test_bad_hook():
3141
packed = packb([3, 1+2j], default=lambda o: o)

0 commit comments

Comments
 (0)
0