diff --git a/extmod/modujson.c b/extmod/modujson.c index 5e90839582..431b5e0cf1 100644 --- a/extmod/modujson.c +++ b/extmod/modujson.c @@ -5,6 +5,8 @@ #include +#include "py/binary.h" +#include "py/objarray.h" #include "py/objlist.h" #include "py/objstringio.h" #include "py/parsenum.h" @@ -53,6 +55,8 @@ typedef struct _ujson_stream_t { mp_obj_t stream_obj; mp_uint_t (*read)(mp_obj_t obj, void *buf, mp_uint_t size, int *errcode); int errcode; + mp_obj_t python_readinto[2 + 1]; + mp_obj_array_t bytearray_obj; byte cur; } ujson_stream_t; @@ -73,9 +77,39 @@ STATIC byte ujson_stream_next(ujson_stream_t *s) { return s->cur; } +STATIC mp_uint_t ujson_python_readinto(mp_obj_t obj, void *buf, mp_uint_t size, int *errcode) { + ujson_stream_t* s = obj; + s->bytearray_obj.items = buf; + s->bytearray_obj.len = size; + *errcode = 0; + mp_obj_t ret = mp_call_method_n_kw(1, 0, s->python_readinto); + if (ret == mp_const_none) { + *errcode = MP_EAGAIN; + return MP_STREAM_ERROR; + } + return mp_obj_get_int(ret); +} + STATIC mp_obj_t _mod_ujson_load(mp_obj_t stream_obj, bool return_first_json) { - const mp_stream_p_t *stream_p = mp_get_stream_raise(stream_obj, MP_STREAM_OP_READ); - ujson_stream_t s = {stream_obj, stream_p->read, 0, 0}; + const mp_stream_p_t *stream_p = mp_proto_get(MP_QSTR_protocol_stream, stream_obj); + ujson_stream_t s; + if (stream_p == NULL) { + mp_load_method(stream_obj, MP_QSTR_readinto, s.python_readinto); + s.bytearray_obj.base.type = &mp_type_bytearray; + s.bytearray_obj.typecode = BYTEARRAY_TYPECODE; + s.bytearray_obj.free = 0; + // len and items are set at read time + s.python_readinto[2] = MP_OBJ_FROM_PTR(&s.bytearray_obj); + s.stream_obj = &s; + s.read = ujson_python_readinto; + } else { + stream_p = mp_get_stream_raise(stream_obj, MP_STREAM_OP_READ); + s.stream_obj = stream_obj; + s.read = stream_p->read; + s.errcode = 0; + s.cur = 0; + } + JSON_DEBUG("got JSON stream\n"); vstr_t vstr; vstr_init(&vstr, 8); diff --git a/tests/extmod/ujson_load_readinto.py b/tests/extmod/ujson_load_readinto.py new file mode 100644 index 0000000000..a277f40efc --- /dev/null +++ b/tests/extmod/ujson_load_readinto.py @@ -0,0 +1,22 @@ +import ujson as json + +# Test that json can load from any object with readinto + +class Buffer: + def __init__(self, data): + self._data = data + self._i = 0 + + def readinto(self, buf): + end = self._i + len(buf) + remaining = len(self._data) - self._i + end = min(end, len(self._data)) + l = min(len(buf), remaining) + buf[:l] = self._data[self._i:end] + self._i += l + return l + +print(json.load(Buffer(b'null'))) +print(json.load(Buffer(b'"abc\\u0064e"'))) +print(json.load(Buffer(b'[false, true, 1, -2]'))) +print(json.load(Buffer(b'{"a":true}'))) diff --git a/tests/extmod/ujson_load_readinto.py.exp b/tests/extmod/ujson_load_readinto.py.exp new file mode 100644 index 0000000000..f8c3c693be --- /dev/null +++ b/tests/extmod/ujson_load_readinto.py.exp @@ -0,0 +1,4 @@ +None +abcde +[False, True, 1, -2] +{'a': True}