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.
This commit is contained in:
Scott Shawcroft 2020-07-10 16:42:32 -07:00
parent 37e77b21cd
commit 734661e79c
No known key found for this signature in database
GPG Key ID: 9349BC7E64B1921E
1 changed files with 34 additions and 2 deletions

View File

@ -26,6 +26,8 @@
#include <stdio.h>
#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);