support json indent

include/pocketpy/pocketpy.h

@@ -616,7 +616,7 @@ PK_API bool py_repr(py_Ref val) PY_RAISE PY_RETURN;
 /// Python equivalent to `len(val)`.
 PK_API bool py_len(py_Ref val) PY_RAISE PY_RETURN;
 /// Python equivalent to `json.dumps(val)`.
-PK_API bool py_json_dumps(py_Ref val) PY_RAISE PY_RETURN;
+PK_API bool py_json_dumps(py_Ref val, int indent) PY_RAISE PY_RETURN;
 /// Python equivalent to `json.loads(val)`.
 PK_API bool py_json_loads(const char* source) PY_RAISE PY_RETURN;
 /// Python equivalent to `pickle.dumps(val)`.

src/modules/json.c

@@ -14,8 +14,10 @@ static bool json_loads(int argc, py_Ref argv) {
 }
 
 static bool json_dumps(int argc, py_Ref argv) {
-    PY_CHECK_ARGC(1);
+    PY_CHECK_ARGC(2);
-    return py_json_dumps(argv);
+    PY_CHECK_ARG_TYPE(1, tp_int);
+    int indent = py_toint(&argv[1]);
+    return py_json_dumps(argv, indent);
 }
 
 void pk__add_module_json() {
@@ -31,47 +33,69 @@ void pk__add_module_json() {
     py_setdict(mod, py_name("Infinity"), &tmp);
 
     py_bindfunc(mod, "loads", json_loads);
-    py_bindfunc(mod, "dumps", json_dumps);
+    py_bind(mod, "dumps(obj, indent=0)", json_dumps);
-}
-
-static bool json__write_object(c11_sbuf* buf, py_TValue* obj);
-
-static bool json__write_array(c11_sbuf* buf, py_TValue* arr, int length) {
-    c11_sbuf__write_char(buf, '[');
-    for(int i = 0; i < length; i++) {
-        if(i != 0) c11_sbuf__write_cstr(buf, ", ");
-        bool ok = json__write_object(buf, arr + i);
-        if(!ok) return false;
-    }
-    c11_sbuf__write_char(buf, ']');
-    return true;
 }
 
 typedef struct {
     c11_sbuf* buf;
     bool first;
+    int indent;
+    int depth;
 } json__write_dict_kv_ctx;
 
+static bool json__write_object(c11_sbuf* buf, py_TValue* obj, int indent, int depth);
+
+static void json__write_indent(c11_sbuf* buf, int n_spaces) {
+    for(int i = 0; i < n_spaces; i++) {
+        c11_sbuf__write_char(buf, ' ');
+    }
+}
+
+static bool json__write_array(c11_sbuf* buf, py_TValue* arr, int length, int indent, int depth) {
+    c11_sbuf__write_char(buf, '[');
+    if(indent > 0) c11_sbuf__write_char(buf, '\n');
+    int n_spaces = indent * depth;
+    const char* sep = indent > 0 ? ",\n" : ", ";
+    for(int i = 0; i < length; i++) {
+        if(i != 0) c11_sbuf__write_cstr(buf, sep);
+        json__write_indent(buf, n_spaces);
+        bool ok = json__write_object(buf, arr + i, indent, depth);
+        if(!ok) return false;
+    }
+    if(indent > 0) {
+        c11_sbuf__write_char(buf, '\n');
+        json__write_indent(buf, n_spaces - indent);
+    }
+    c11_sbuf__write_char(buf, ']');
+    return true;
+}
+
 static bool json__write_dict_kv(py_Ref k, py_Ref v, void* ctx_) {
     json__write_dict_kv_ctx* ctx = ctx_;
-    if(!ctx->first) c11_sbuf__write_cstr(ctx->buf, ", ");
+    int n_spaces = ctx->indent * ctx->depth;
+    const char* sep = ctx->indent > 0 ? ",\n" : ", ";
+    if(!ctx->first) c11_sbuf__write_cstr(ctx->buf, sep);
     ctx->first = false;
     if(!py_isstr(k)) return TypeError("keys must be strings");
+    json__write_indent(ctx->buf, n_spaces);
     c11_sbuf__write_quoted(ctx->buf, py_tosv(k), '"');
     c11_sbuf__write_cstr(ctx->buf, ": ");
-    return json__write_object(ctx->buf, v);
+    return json__write_object(ctx->buf, v, ctx->indent, ctx->depth);
 }
 
 static bool json__write_namedict_kv(py_Name k, py_Ref v, void* ctx_) {
     json__write_dict_kv_ctx* ctx = ctx_;
-    if(!ctx->first) c11_sbuf__write_cstr(ctx->buf, ", ");
+    int n_spaces = ctx->indent * ctx->depth;
+    const char* sep = ctx->indent > 0 ? ",\n" : ", ";
+    if(!ctx->first) c11_sbuf__write_cstr(ctx->buf, sep);
     ctx->first = false;
+    json__write_indent(ctx->buf, n_spaces);
     c11_sbuf__write_quoted(ctx->buf, py_name2sv(k), '"');
     c11_sbuf__write_cstr(ctx->buf, ": ");
-    return json__write_object(ctx->buf, v);
+    return json__write_object(ctx->buf, v, ctx->indent, ctx->depth);
 }
 
-static bool json__write_object(c11_sbuf* buf, py_TValue* obj) {
+static bool json__write_object(c11_sbuf* buf, py_TValue* obj, int indent, int depth) {
     switch(obj->type) {
         case tp_NoneType: c11_sbuf__write_cstr(buf, "null"); return true;
         case tp_int: c11_sbuf__write_int(buf, obj->_i64); return true;
@@ -94,24 +118,40 @@ static bool json__write_object(c11_sbuf* buf, py_TValue* obj) {
             return true;
         }
         case tp_list: {
-            return json__write_array(buf, py_list_data(obj), py_list_len(obj));
+            return json__write_array(buf, py_list_data(obj), py_list_len(obj), indent, depth + 1);
         }
         case tp_tuple: {
-            return json__write_array(buf, py_tuple_data(obj), py_tuple_len(obj));
+            return json__write_array(buf, py_tuple_data(obj), py_tuple_len(obj), indent, depth + 1);
         }
         case tp_dict: {
             c11_sbuf__write_char(buf, '{');
-            json__write_dict_kv_ctx ctx = {.buf = buf, .first = true};
+            if(indent > 0) c11_sbuf__write_char(buf, '\n');
+            json__write_dict_kv_ctx ctx = {.buf = buf,
+                                           .first = true,
+                                           .indent = indent,
+                                           .depth = depth + 1};
             bool ok = py_dict_apply(obj, json__write_dict_kv, &ctx);
             if(!ok) return false;
+            if(indent > 0) {
+                c11_sbuf__write_char(buf, '\n');
+                json__write_indent(buf, indent * depth);
+            }
             c11_sbuf__write_char(buf, '}');
             return true;
         }
         case tp_namedict: {
             c11_sbuf__write_char(buf, '{');
-            json__write_dict_kv_ctx ctx = {.buf = buf, .first = true};
+            if(indent > 0) c11_sbuf__write_char(buf, '\n');
+            json__write_dict_kv_ctx ctx = {.buf = buf,
+                                           .first = true,
+                                           .indent = indent,
+                                           .depth = depth + 1};
             bool ok = py_applydict(py_getslot(obj, 0), json__write_namedict_kv, &ctx);
             if(!ok) return false;
+            if(indent > 0) {
+                c11_sbuf__write_char(buf, '\n');
+                json__write_indent(buf, indent * depth);
+            }
             c11_sbuf__write_char(buf, '}');
             return true;
         }
@@ -119,10 +159,10 @@ static bool json__write_object(c11_sbuf* buf, py_TValue* obj) {
     }
 }
 
-bool py_json_dumps(py_Ref val) {
+bool py_json_dumps(py_Ref val, int indent) {
     c11_sbuf buf;
     c11_sbuf__ctor(&buf);
-    bool ok = json__write_object(&buf, val);
+    bool ok = json__write_object(&buf, val, indent, 0);
     if(!ok) {
         c11_sbuf__dtor(&buf);
         return false;

tests/73_json_indent.py

@@ -0,0 +1,47 @@
+import json
+
+assert json.dumps([1, 2, [3, 4], 5], indent=2) == '[\n  1,\n  2,\n  [\n    3,\n    4\n  ],\n  5\n]'
+
+a = {
+    'a': 1,
+    'b': 2,
+    'c': None,
+    'd': [1, 2, 3],
+    'e': {
+        'a': 100,
+        'b': 2.5,
+        'c': None,
+        'd': [142, 2785, 39767],
+    },
+    "f": 'This is a string',
+    'g': [True, False, None],
+    'h': False
+}
+
+assert json.dumps(a, indent=2) == '''{
+  "a": 1,
+  "b": 2,
+  "c": null,
+  "d": [
+    1,
+    2,
+    3
+  ],
+  "e": {
+    "a": 100,
+    "b": 2.5,
+    "c": null,
+    "d": [
+      142,
+      2785,
+      39767
+    ]
+  },
+  "f": "This is a string",
+  "g": [
+    true,
+    false,
+    null
+  ],
+  "h": false
+}'''