diff --git a/lib/json/json.go b/lib/json/json.go index bc67be5d..eb16b2c9 100644 --- a/lib/json/json.go +++ b/lib/json/json.go @@ -28,6 +28,7 @@ import ( // // json = module( // encode, +// encode_indent, // decode, // indent, // ) @@ -54,6 +55,10 @@ import ( // (e.g. it implements both Iterable and HasFields), the first case takes precedence. // Encoding any other value yields an error. // +// def encode_indent(x, *, prefix="", indent="\t"): +// +// Equivalent to json.indent(json.encode(x), prefix=prefix, indent=indent). +// // def decode(x[, default]): // // The decode function has one required positional parameter, a JSON string. @@ -76,9 +81,10 @@ import ( var Module = &starlarkstruct.Module{ Name: "json", Members: starlark.StringDict{ - "encode": starlark.NewBuiltin("json.encode", encode), - "decode": starlark.NewBuiltin("json.decode", decode), - "indent": starlark.NewBuiltin("json.indent", indent), + "encode": starlark.NewBuiltin("json.encode", encode), + "encode_indent": starlark.NewBuiltin("json.encode_indent", encodeIndent), + "decode": starlark.NewBuiltin("json.decode", decode), + "indent": starlark.NewBuiltin("json.indent", indent), }, } @@ -233,6 +239,27 @@ func encode(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, k return starlark.String(buf.String()), nil } +func encodeIndent(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + prefix, indent := "", "\t" // keyword-only + if err := starlark.UnpackArgs(b.Name(), nil, kwargs, + "prefix?", &prefix, + "indent?", &indent, + ); err != nil { + return nil, err + } + // Rely on encode() to parse the positional parameter (since the signature matches); use our b so + // error messages are attributed to json.encode_indent. + str, err := encode(thread, b, args, nil) + if err != nil { + return nil, err + } + var buf bytes.Buffer + if err := json.Indent(&buf, []byte(str.(starlark.String)), prefix, indent); err != nil { + return nil, fmt.Errorf("%s: %v", b.Name(), err) + } + return starlark.String(buf.String()), nil +} + func pointer(i interface{}) unsafe.Pointer { v := reflect.ValueOf(i) switch v.Kind() { diff --git a/starlark/testdata/json.star b/starlark/testdata/json.star index 1df9dee9..47ae0eb6 100644 --- a/starlark/testdata/json.star +++ b/starlark/testdata/json.star @@ -3,7 +3,7 @@ load("assert.star", "assert") load("json.star", "json") -assert.eq(dir(json), ["decode", "encode", "indent"]) +assert.eq(dir(json), ["decode", "encode", "encode_indent", "indent"]) # Some of these cases were inspired by github.com/nst/JSONTestSuite. @@ -49,6 +49,20 @@ recursive_tuple = (1, 2, []) recursive_tuple[2].append(recursive_tuple) encode_error(recursive_tuple, 'json.encode: at tuple index 2: at list index 0: cycle in JSON structure') +## json.encode_indent + +assert.eq( + json.encode_indent({"x": 1, "y": ["one", "two"]}, prefix="¶", indent="–––"), + json.indent(json.encode({"x": 1, "y": ["one", "two"]}), prefix="¶", indent="–––") +) +assert.eq( + json.encode_indent({"x": 1, "y": ["one", "two"]}), + json.indent(json.encode({"x": 1, "y": ["one", "two"]}), prefix="", indent="\t") +) + +assert.fails(lambda: json.encode_indent('foo', 'bar'), 'json.encode_indent: got 2 arguments, want 1') +assert.fails(lambda: json.encode_indent(recursive_map), 'json.encode_indent: in dict key "r": cycle in JSON structure') + ## json.decode assert.eq(json.decode("null"), None)