From 734661e79cc2b108079377642ef24fab1217c484 Mon Sep 17 00:00:00 2001 From: Scott Shawcroft Date: Fri, 10 Jul 2020 16:42:32 -0700 Subject: [PATCH 1/2] Add support to json.load for any object with readinto This way we don't need to load the whole string version of the json into memory. --- extmod/modujson.c | 36 ++++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/extmod/modujson.c b/extmod/modujson.c index 0f93ccb110..1e831783a8 100644 --- a/extmod/modujson.c +++ b/extmod/modujson.c @@ -26,6 +26,8 @@ #include +#include "py/binary.h" +#include "py/objarray.h" #include "py/objlist.h" #include "py/objstringio.h" #include "py/parsenum.h" @@ -74,6 +76,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; @@ -94,9 +98,37 @@ 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; + } + JSON_DEBUG("got JSON stream\n"); vstr_t vstr; vstr_init(&vstr, 8); From 372bcf8a956c8b20c423333114d7aa061dd54751 Mon Sep 17 00:00:00 2001 From: Scott Shawcroft Date: Fri, 10 Jul 2020 17:33:17 -0700 Subject: [PATCH 2/2] Fix stream version and add basic readinto test --- extmod/modujson.c | 2 ++ tests/extmod/ujson_load_readinto.py | 22 ++++++++++++++++++++++ tests/extmod/ujson_load_readinto.py.exp | 4 ++++ 3 files changed, 28 insertions(+) create mode 100644 tests/extmod/ujson_load_readinto.py create mode 100644 tests/extmod/ujson_load_readinto.py.exp diff --git a/extmod/modujson.c b/extmod/modujson.c index 1e831783a8..242726cca0 100644 --- a/extmod/modujson.c +++ b/extmod/modujson.c @@ -127,6 +127,8 @@ STATIC mp_obj_t _mod_ujson_load(mp_obj_t stream_obj, bool return_first_json) { 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"); 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}