extmod/modujson: Support passing bytes/bytearray to json.loads.

CPython allows this, and it can be useful to reduce the number of memory
allocations.

Fixes issue #5031.

2 files changed, 8 insertions, 4 deletions

diff --git a/extmod/modujson.c b/extmod/modujson.c
index f5c6428ef..15ed2f38d 100644
--- a/extmod/modujson.c
+++ b/extmod/modujson.c
@@ -3,7 +3,7 @@
  *
  * The MIT License (MIT)
  *
- * Copyright (c) 2014-2016 Damien P. George
+ * Copyright (c) 2014-2019 Damien P. George
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -281,9 +281,9 @@ STATIC mp_obj_t mod_ujson_load(mp_obj_t stream_obj) {
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_ujson_load_obj, mod_ujson_load);
 
 STATIC mp_obj_t mod_ujson_loads(mp_obj_t obj) {
-    size_t len;
-    const char *buf = mp_obj_str_get_data(obj, &len);
-    vstr_t vstr = {len, len, (char*)buf, true};
+    mp_buffer_info_t bufinfo;
+    mp_get_buffer_raise(obj, &bufinfo, MP_BUFFER_READ);
+    vstr_t vstr = {bufinfo.len, bufinfo.len, (char*)bufinfo.buf, true};
     mp_obj_stringio_t sio = {{&mp_type_stringio}, &vstr, 0, MP_OBJ_NULL};
     return mod_ujson_load(MP_OBJ_FROM_PTR(&sio));
 }
diff --git a/tests/extmod/ujson_loads.py b/tests/extmod/ujson_loads.py
index adba3c068..43672d650 100644
--- a/tests/extmod/ujson_loads.py
+++ b/tests/extmod/ujson_loads.py
@@ -37,6 +37,10 @@ my_print(json.loads('"abc\\uabcd"'))
 # whitespace handling
 my_print(json.loads('{\n\t"a":[]\r\n, "b":[1], "c":{"3":4}     \n\r\t\r\r\r\n}'))
 
+# loading from bytes and bytearray
+my_print(json.loads(b'[1,2]'))
+my_print(json.loads(bytearray(b'[null]')))
+
 # loading nothing should raise exception
 try:
     json.loads('')