Skip to content

Commit 734661e

Browse files
committed
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.
1 parent 37e77b2 commit 734661e

File tree

1 file changed

+34
-2
lines changed

1 file changed

+34
-2
lines changed

extmod/modujson.c

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626

2727
#include <stdio.h>
2828

29+
#include "py/binary.h"
30+
#include "py/objarray.h"
2931
#include "py/objlist.h"
3032
#include "py/objstringio.h"
3133
#include "py/parsenum.h"
@@ -74,6 +76,8 @@ typedef struct _ujson_stream_t {
7476
mp_obj_t stream_obj;
7577
mp_uint_t (*read)(mp_obj_t obj, void *buf, mp_uint_t size, int *errcode);
7678
int errcode;
79+
mp_obj_t python_readinto[2 + 1];
80+
mp_obj_array_t bytearray_obj;
7781
byte cur;
7882
} ujson_stream_t;
7983

@@ -94,9 +98,37 @@ STATIC byte ujson_stream_next(ujson_stream_t *s) {
9498
return s->cur;
9599
}
96100

101+
STATIC mp_uint_t ujson_python_readinto(mp_obj_t obj, void *buf, mp_uint_t size, int *errcode) {
102+
ujson_stream_t* s = obj;
103+
s->bytearray_obj.items = buf;
104+
s->bytearray_obj.len = size;
105+
*errcode = 0;
106+
mp_obj_t ret = mp_call_method_n_kw(1, 0, s->python_readinto);
107+
if (ret == mp_const_none) {
108+
*errcode = MP_EAGAIN;
109+
return MP_STREAM_ERROR;
110+
}
111+
return mp_obj_get_int(ret);
112+
}
113+
97114
STATIC mp_obj_t _mod_ujson_load(mp_obj_t stream_obj, bool return_first_json) {
98-
const mp_stream_p_t *stream_p = mp_get_stream_raise(stream_obj, MP_STREAM_OP_READ);
99-
ujson_stream_t s = {stream_obj, stream_p->read, 0, 0};
115+
const mp_stream_p_t *stream_p = mp_proto_get(MP_QSTR_protocol_stream, stream_obj);
116+
ujson_stream_t s;
117+
if (stream_p == NULL) {
118+
mp_load_method(stream_obj, MP_QSTR_readinto, s.python_readinto);
119+
s.bytearray_obj.base.type = &mp_type_bytearray;
120+
s.bytearray_obj.typecode = BYTEARRAY_TYPECODE;
121+
s.bytearray_obj.free = 0;
122+
// len and items are set at read time
123+
s.python_readinto[2] = MP_OBJ_FROM_PTR(&s.bytearray_obj);
124+
s.stream_obj = &s;
125+
s.read = ujson_python_readinto;
126+
} else {
127+
stream_p = mp_get_stream_raise(stream_obj, MP_STREAM_OP_READ);
128+
s.stream_obj = stream_obj;
129+
s.read = stream_p->read;
130+
}
131+
100132
JSON_DEBUG("got JSON stream\n");
101133
vstr_t vstr;
102134
vstr_init(&vstr, 8);

0 commit comments

Comments
 (0)