8000 Implement object_pairs_hook · twigtek/msgpack-python@7794251 · GitHub
[go: up one dir, main page]
Skip to content

Commit 7794251

Browse files
jnothmanjnothman
committed
Implement object_pairs_hook
1 parent b06ed8e commit 7794251

File tree

5 files changed

+61
-22
lines changed

5 files changed

+61
-22
lines changed

‎msgpack/_msgpack.pyx

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,7 @@ cdef extern from "unpack.h":
197197
ctypedef struct msgpack_user:
198198
int use_list
199199
PyObject* object_hook
200+
bint has_pairs_hook # call object_hook with k-v pairs
200201
PyObject* list_hook
201202
char *encoding
202203
char *unicode_errors
@@ -213,18 +214,32 @@ cdef extern from "unpack.h":
213214
void template_init(template_context* ctx)
214215
object template_data(template_context* ctx)
215216

216-
cdef inline init_ctx(template_context *ctx, object object_hook, object list_hook, bint use_list, encoding, unicode_errors):
217+
cdef inline init_ctx(template_context *ctx, object object_hook, object object_pairs_hook, object list_hook, bint use_list, encoding, unicode_errors):
217218
template_init(ctx)
218219
ctx.user.use_list = use_list
219220
ctx.user.object_hook = ctx.user.list_hook = <PyObject*>NULL
221+
222+
if object_hook is not None and object_pairs_hook is not None:
223+
raise ValueError("object_pairs_hook and object_hook are mutually exclusive.")
224+
220225
if object_hook is not None:
221226
if not PyCallable_Check(object_hook):
222227
raise TypeError("object_hook must be a callable.")
223228
ctx.user.object_hook = <PyObject*>object_hook
229+
230+
if object_pairs_hook is None:
231+
ctx.user.has_pairs_hook = False
232+
else:
233+
if not PyCallable_Check(object_pairs_hook):
234+
raise TypeError("object_pairs_hook must be a callable.")
235+
ctx.user.object_hook = <PyObject*>object_pairs_hook
236+
ctx.user.has_pairs_hook = True
237+
224238
if list_hook is not None:
225239
if not PyCallable_Check(list_hook):
226240
raise TypeError("list_hook must be a callable.")
227241
ctx.user.list_hook = <PyObject*>list_hook
242+
228243
if encoding is None:
229244
ctx.user.encoding = NULL
230245
ctx.user.unicode_errors = NULL
@@ -240,7 +255,7 @@ cdef inline init_ctx(template_context *ctx, object object_hook, object list_hook
240255
_berrors = unicode_errors
241256
ctx.user.unicode_errors = PyBytes_AsString(_berrors)
242257

243-
def unpackb(object packed, object object_hook=None, object list_hook=None,
258+
def unpackb(object packed, object object_hook=None, object object_pairs_hook=None, object list_hook=None,
244259
bint use_list=0, encoding=None, unicode_errors="strict",
245260
):
246261
"""Unpack packed_bytes to object. Returns an unpacked object.
@@ -255,7 +270,7 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
255270
cdef Py_ssize_t buf_len
256271
PyObject_AsReadBuffer(packed, <const_void_ptr*>&buf, &buf_len)
257272

258-
init_ctx(&ctx, object_hook, list_hook, use_list, encoding, unicode_errors)
273+
init_ctx(&ctx, object_hook, object_pairs_hook, list_hook, use_list, encoding, unicode_errors)
259274
ret = template_execute(&ctx, buf, buf_len, &off, 1)
260275
if ret == 1:
261276
obj = template_data(&ctx)
@@ -266,15 +281,15 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
266281
return None
267282

268283

269-
def unpack(object stream, object object_hook=None, object list_hook=None,
284+
def unpack(object stream, object object_hook=None, object object_pairs_hook=None, object list_hook=None,
270285
bint use_list=0, encoding=None, unicode_errors="strict",
271286
):
272287
"""Unpack an object from `stream`.
273288
274289
Raises `ValueError` when `stream` has extra bytes.
275290
"""
276291
return unpackb(stream.read(), use_list=use_list,
277-
object_hook=object_hook, list_hook=list_hook,
292+
object_hook=object_hook, object_pairs_hook=object_pairs_hook, list_hook=list_hook,
278293
encoding=encoding, unicode_errors=unicode_errors,
279294
)
280295

@@ -294,7 +309,10 @@ cdef class Unpacker(object):
294309
Otherwise, it is deserialized to Python tuple. (default: False)
295310
296311
`object_hook` is same to simplejson. If it is not None, it should be callable
297-
and Unpacker calls it when deserializing key-value.
312+
and Unpacker calls it with a dict argument after deserializing a map.
313+
314+
`object_pairs_hook` is same to simplejson. If it is not None, it should be callable
315+
and Unpacker calls it with a list of key-value pairs after deserializing a map.
298316
299317
`encoding` is encoding used for decoding msgpack bytes. If it is None (default),
300318
msgpack bytes is deserialized to Python bytes.
@@ -345,7 +363,7 @@ cdef class Unpacker(object):
345363
self.buf = NULL
346364

347365
def __init__(self, file_like=None, Py_ssize_t read_size=0, bint use_list=0,
348-
object object_hook=None, object list_hook=None,
366+
object object_hook=None, object object_pairs_hook=None, object list_hook=None,
349367
encoding=None, unicode_errors='strict', int max_buffer_size=0,
350368
):
351369
self.use_list = use_list
@@ -368,7 +386,7 @@ cdef class Unpacker(object):
368386
self.buf_size = read_size
369387
self.buf_head = 0
370388
self.buf_tail = 0
371-
init_ctx(&self.ctx, object_hook, list_hook, use_list, encoding, unicode_errors)
389+
init_ctx(&self.ctx, object_hook, object_pairs_hook, list_hook, use_list, encoding, unicode_errors)
372390

373391
def feed(self, object next_bytes):
374392
cdef char* buf

‎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: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -357,7 +357,7 @@ msgpack_unpack_func(int, _execute)(msgpack_unpack_struct(_context)* ctx, const c
357357
c->ct = CT_MAP_VALUE;
358358
goto _header_again;
359359
case CT_MAP_VALUE:
360-
if(construct_cb(_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; }
361361
if(++c->count == c->size) {
362362
obj = c->obj;
363363
construct_cb(_map_end)(user, &obj);

‎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)
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('', 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)

‎test/test_pack.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -111,10 +111,9 @@ def test_odict():
111111
seq = [(b'one', 1), (b'two', 2), (b'three', 3), (b'four', 4)]
112112
od = odict(seq)
113113
assert_equal(unpackb(packb(od)), dict(seq))
114-
# After object_pairs_hook is implemented.
115-
#def pair_hook(seq):
116-
# return seq
117-
#assert_equal(unpackb(packb(od), object_pairs_hook=pair_hook), seq)
114+
def pair_hook(seq):
115+
return seq
116+
assert_equal(unpackb(packb(od), object_pairs_hook=pair_hook), seq)
118117

119118

120119
if __name__ == '__main__':

0 commit comments

Comments
 (0)
0