extmod/modujson: Add support for dump/dumps separators keyword-argument.

Optionally enabled via MICROPY_PY_UJSON_SEPARATORS.  Enabled by default.

For dump, make sure mp_get_stream_raise is called after
mod_ujson_separators since CPython does it in this order (if both
separators and stream are invalid, separators will raise an exception
first).

Add separators argument in the docs as well.

Signed-off-by: Peter Züger <zueger.peter@icloud.com>
Signed-off-by: Damien George <damien@micropython.org>

1 files changed, 58 insertions, 0 deletions

diff --git a/extmod/modujson.c b/extmod/modujson.c
index 8dff67358..e5bfb0d96 100644
--- a/extmod/modujson.c
+++ b/extmod/modujson.c
@@ -34,6 +34,62 @@
 
 #if MICROPY_PY_UJSON
 
+#if MICROPY_PY_UJSON_SEPARATORS
+
+enum {
+    DUMP_MODE_TO_STRING = 1,
+    DUMP_MODE_TO_STREAM = 2,
+};
+
+STATIC mp_obj_t mod_ujson_dump_helper(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args, unsigned int mode) {
+    enum { ARG_separators };
+    static const mp_arg_t allowed_args[] = {
+        { MP_QSTR_separators, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_rom_obj = MP_ROM_NONE} },
+    };
+
+    mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
+    mp_arg_parse_all(n_args - mode, pos_args + mode, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
+
+    mp_print_ext_t print_ext;
+
+    if (args[ARG_separators].u_obj == mp_const_none) {
+        print_ext.item_separator = ", ";
+        print_ext.key_separator = ": ";
+    } else {
+        mp_obj_t *items;
+        mp_obj_get_array_fixed_n(args[ARG_separators].u_obj, 2, &items);
+        print_ext.item_separator = mp_obj_str_get_str(items[0]);
+        print_ext.key_separator = mp_obj_str_get_str(items[1]);
+    }
+
+    if (mode == DUMP_MODE_TO_STRING) {
+        // dumps(obj)
+        vstr_t vstr;
+        vstr_init_print(&vstr, 8, &print_ext.base);
+        mp_obj_print_helper(&print_ext.base, pos_args[0], PRINT_JSON);
+        return mp_obj_new_str_from_vstr(&mp_type_str, &vstr);
+    } else {
+        // dump(obj, stream)
+        print_ext.base.data = MP_OBJ_TO_PTR(pos_args[1]);
+        print_ext.base.print_strn = mp_stream_write_adaptor;
+        mp_get_stream_raise(pos_args[1], MP_STREAM_OP_WRITE);
+        mp_obj_print_helper(&print_ext.base, pos_args[0], PRINT_JSON);
+        return mp_const_none;
+    }
+}
+
+STATIC mp_obj_t mod_ujson_dump(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
+    return mod_ujson_dump_helper(n_args, pos_args, kw_args, DUMP_MODE_TO_STREAM);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_KW(mod_ujson_dump_obj, 2, mod_ujson_dump);
+
+STATIC mp_obj_t mod_ujson_dumps(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
+    return mod_ujson_dump_helper(n_args, pos_args, kw_args, DUMP_MODE_TO_STRING);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_KW(mod_ujson_dumps_obj, 1, mod_ujson_dumps);
+
+#else
+
 STATIC mp_obj_t mod_ujson_dump(mp_obj_t obj, mp_obj_t stream) {
     mp_get_stream_raise(stream, MP_STREAM_OP_WRITE);
     mp_print_t print = {MP_OBJ_TO_PTR(stream), mp_stream_write_adaptor};
@@ -51,6 +107,8 @@ STATIC mp_obj_t mod_ujson_dumps(mp_obj_t obj) {
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_ujson_dumps_obj, mod_ujson_dumps);
 
+#endif
+
 // The function below implements a simple non-recursive JSON parser.
 //
 // The JSON specification is at http://www.ietf.org/rfc/rfc4627.txt