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>

5 files changed, 33 insertions, 4 deletions

diff --git a/py/mpconfig.h b/py/mpconfig.h
index f67e11cd4..a91c39b01 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -1380,6 +1380,11 @@ typedef double mp_float_t;
 #define MICROPY_PY_UJSON (0)
 #endif
 
+// Whether to support the "separators" argument to dump, dumps
+#ifndef MICROPY_PY_UJSON_SEPARATORS
+#define MICROPY_PY_UJSON_SEPARATORS (1)
+#endif
+
 #ifndef MICROPY_PY_URE
 #define MICROPY_PY_URE (0)
 #endif
diff --git a/py/mpprint.h b/py/mpprint.h
index aef9015de..0dff9a770 100644
--- a/py/mpprint.h
+++ b/py/mpprint.h
@@ -52,6 +52,14 @@ typedef struct _mp_print_t {
     mp_print_strn_t print_strn;
 } mp_print_t;
 
+typedef struct _mp_print_ext_t {
+    mp_print_t base;
+    const char *item_separator;
+    const char *key_separator;
+}mp_print_ext_t;
+
+#define MP_PRINT_GET_EXT(print) ((mp_print_ext_t *)print)
+
 // All (non-debug) prints go through one of the two interfaces below.
 // 1) Wrapper for platform print function, which wraps MP_PLAT_PRINT_STRN.
 extern const mp_print_t mp_plat_print;
diff --git a/py/objdict.c b/py/objdict.c
index 63e5381c6..ed4376aa4 100644
--- a/py/objdict.c
+++ b/py/objdict.c
@@ -69,8 +69,15 @@ STATIC mp_map_elem_t *dict_iter_next(mp_obj_dict_t *dict, size_t *cur) {
 STATIC void dict_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind) {
     mp_obj_dict_t *self = MP_OBJ_TO_PTR(self_in);
     bool first = true;
+    const char *item_separator = ", ";
+    const char *key_separator = ": ";
     if (!(MICROPY_PY_UJSON && kind == PRINT_JSON)) {
         kind = PRINT_REPR;
+    } else {
+        #if MICROPY_PY_UJSON_SEPARATORS
+        item_separator = MP_PRINT_GET_EXT(print)->item_separator;
+        key_separator = MP_PRINT_GET_EXT(print)->key_separator;
+        #endif
     }
     if (MICROPY_PY_COLLECTIONS_ORDEREDDICT && self->base.type != &mp_type_dict && kind != PRINT_JSON) {
         mp_printf(print, "%q(", self->base.type->name);
@@ -80,7 +87,7 @@ STATIC void dict_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_
     mp_map_elem_t *next = NULL;
     while ((next = dict_iter_next(self, &cur)) != NULL) {
         if (!first) {
-            mp_print_str(print, ", ");
+            mp_print_str(print, item_separator);
         }
         first = false;
         bool add_quote = MICROPY_PY_UJSON && kind == PRINT_JSON && !mp_obj_is_str_or_bytes(next->key);
@@ -91,7 +98,7 @@ STATIC void dict_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_
         if (add_quote) {
             mp_print_str(print, "\"");
         }
-        mp_print_str(print, ": ");
+        mp_print_str(print, key_separator);
         mp_obj_print_helper(print, next->value, kind);
     }
     mp_print_str(print, "}");
diff --git a/py/objlist.c b/py/objlist.c
index 8c989facc..f431e273d 100644
--- a/py/objlist.c
+++ b/py/objlist.c
@@ -44,13 +44,18 @@ STATIC mp_obj_t list_pop(size_t n_args, const mp_obj_t *args);
 
 STATIC void list_print(const mp_print_t *print, mp_obj_t o_in, mp_print_kind_t kind) {
     mp_obj_list_t *o = MP_OBJ_TO_PTR(o_in);
+    const char *item_separator = ", ";
     if (!(MICROPY_PY_UJSON && kind == PRINT_JSON)) {
         kind = PRINT_REPR;
+    } else {
+        #if MICROPY_PY_UJSON_SEPARATORS
+        item_separator = MP_PRINT_GET_EXT(print)->item_separator;
+        #endif
     }
     mp_print_str(print, "[");
     for (size_t i = 0; i < o->len; i++) {
         if (i > 0) {
-            mp_print_str(print, ", ");
+            mp_print_str(print, item_separator);
         }
         mp_obj_print_helper(print, o->items[i], kind);
     }
diff --git a/py/objtuple.c b/py/objtuple.c
index 07a560ac0..67d7bc356 100644
--- a/py/objtuple.c
+++ b/py/objtuple.c
@@ -39,15 +39,19 @@
 
 void mp_obj_tuple_print(const mp_print_t *print, mp_obj_t o_in, mp_print_kind_t kind) {
     mp_obj_tuple_t *o = MP_OBJ_TO_PTR(o_in);
+    const char *item_separator = ", ";
     if (MICROPY_PY_UJSON && kind == PRINT_JSON) {
         mp_print_str(print, "[");
+        #if MICROPY_PY_UJSON_SEPARATORS
+        item_separator = MP_PRINT_GET_EXT(print)->item_separator;
+        #endif
     } else {
         mp_print_str(print, "(");
         kind = PRINT_REPR;
     }
     for (size_t i = 0; i < o->len; i++) {
         if (i > 0) {
-            mp_print_str(print, ", ");
+            mp_print_str(print, item_separator);
         }
         mp_obj_print_helper(print, o->items[i], kind);
     }