Merge pull request #3143 from tannewt/improve_json

Add support to json.load for any object with readinto

3 files changed, 62 insertions, 2 deletions

diff --git a/extmod/modujson.c b/extmod/modujson.c
index 5e9083958..431b5e0cf 100644
--- a/extmod/modujson.c
+++ b/extmod/modujson.c
@@ -5,6 +5,8 @@
 
 #include <stdio.h>
 
+#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 000000000..a277f40ef
--- /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 000000000..f8c3c693b
--- /dev/null
+++ b/tests/extmod/ujson_load_readinto.py.exp
@@ -0,0 +1,4 @@
+None
+abcde
+[False, True, 1, -2]
+{'a': True}