From 5c56e846db9f452e756b51c1f0dbddb298f7e463 Mon Sep 17 00:00:00 2001 From: Brian Grenier Date: Wed, 30 Oct 2024 14:11:11 -0600 Subject: [PATCH] fix nested pretty print --- ember_json/array.mojo | 14 ++++++++------ ember_json/json.mojo | 2 +- ember_json/object.mojo | 11 +++++++---- ember_json/traits.mojo | 4 +--- ember_json/utils.mojo | 7 +++++-- ember_json/value.mojo | 12 ++++++++---- test/ember_json/test_json.mojo | 23 +++++++++++++++++++++++ 7 files changed, 53 insertions(+), 20 deletions(-) diff --git a/ember_json/array.mojo b/ember_json/array.mojo index f466ecd..4f9b2e3 100644 --- a/ember_json/array.mojo +++ b/ember_json/array.mojo @@ -4,7 +4,7 @@ from .reader import Reader from .utils import * from .constants import * from sys.intrinsics import unlikely, likely -from .traits import JsonValue, PrettyPrintable, DefaultPrettyIndent +from .traits import JsonValue, PrettyPrintable @value @@ -58,17 +58,19 @@ struct Array(Sized, JsonValue): (",").write_to(writer) ("]").write_to(writer) - fn pretty_to[W: Writer](self, inout writer: W, indent: Variant[Int, String] = DefaultPrettyIndent): - var ind = String(" ") * indent[Int] if indent.isa[Int]() else indent[String] + fn pretty_to[W: Writer](self, inout writer: W, indent: String): writer.write("[\n") + self._pretty_write_items(writer, indent) + writer.write("]") + + fn _pretty_write_items[W: Writer](self, inout writer: W, indent: String): for i in range(len(self._data)): - writer.write(ind, self._data[i]) + writer.write(indent) + self._data[i].pretty_to(writer, indent) if i < len(self._data) - 1: writer.write(",") writer.write("\n") - writer.write("]") - @always_inline fn __str__(self) -> String: return write(self) diff --git a/ember_json/json.mojo b/ember_json/json.mojo index f4a946b..40b0ec0 100644 --- a/ember_json/json.mojo +++ b/ember_json/json.mojo @@ -138,7 +138,7 @@ struct JSON(JsonValue, Sized, PrettyPrintable): else: self.array().write_to(writer) - fn pretty_to[W: Writer](self, inout writer: W, indent: Variant[Int, String] = DefaultPrettyIndent): + fn pretty_to[W: Writer](self, inout writer: W, indent: String): """Write the pretty representation to a writer. Args: diff --git a/ember_json/object.mojo b/ember_json/object.mojo index 2fb0770..fd54b64 100644 --- a/ember_json/object.mojo +++ b/ember_json/object.mojo @@ -78,20 +78,23 @@ struct Object(Sized, JsonValue, PrettyPrintable): ("}").write_to(writer) - fn pretty_to[W: Writer](self, inout writer: W, indent: Variant[Int, String] = DefaultPrettyIndent): - var ind = String(" ") * indent[Int] if indent.isa[Int]() else indent[String] + fn pretty_to[W: Writer](self, inout writer: W, indent: String): writer.write("{\n") + self._pretty_write_items(writer, indent) + writer.write("}") + + fn _pretty_write_items[W: Writer](self, inout writer: W, indent: String): try: var done = 0 for k in self._data: - writer.write(ind, '"', k[], '"', ": ", self[k[]]) + writer.write(indent, '"', k[], '"', ": ") + self[k[]].pretty_to(writer, indent) if done < len(self._data) - 1: writer.write(",") writer.write("\n") done += 1 except: pass - writer.write("}") @always_inline fn __str__(self) -> String: diff --git a/ember_json/traits.mojo b/ember_json/traits.mojo index 287319b..68344ce 100644 --- a/ember_json/traits.mojo +++ b/ember_json/traits.mojo @@ -1,10 +1,8 @@ from utils import Variant -alias DefaultPrettyIndent = String(" ") * 4 - trait PrettyPrintable: - fn pretty_to[W: Writer](self, inout writer: W, indent: Variant[Int, String] = DefaultPrettyIndent): + fn pretty_to[W: Writer](self, inout writer: W, indent: String): ... diff --git a/ember_json/utils.mojo b/ember_json/utils.mojo index 5b6bfca..add2882 100644 --- a/ember_json/utils.mojo +++ b/ember_json/utils.mojo @@ -2,12 +2,14 @@ from .constants import * from utils import Span, Variant from memory import memcmp, memcpy from utils.write import _WriteBuffer -from .traits import JsonValue, PrettyPrintable, DefaultPrettyIndent +from .traits import JsonValue, PrettyPrintable alias Bytes = String._buffer_type alias ByteVec = SIMD[DType.uint8, _] alias ByteView = Span[Byte, _] +alias DefaultPrettyIndent = 4 + @always_inline fn write[T: JsonValue, //](v: T) -> String: @@ -19,7 +21,8 @@ fn write[T: JsonValue, //](v: T) -> String: fn write_pretty[P: PrettyPrintable, //](v: P, indent: Variant[Int, String] = DefaultPrettyIndent) -> String: var writer = _WriteBuffer[4096](String()) - v.pretty_to(writer, indent) + var ind = String(" ") * indent[Int] if indent.isa[Int]() else indent[String] + v.pretty_to(writer, ind) writer.flush() return writer.writer diff --git a/ember_json/value.mojo b/ember_json/value.mojo index 95f5baa..04ced25 100644 --- a/ember_json/value.mojo +++ b/ember_json/value.mojo @@ -249,11 +249,15 @@ struct Value(JsonValue): elif self.isa[Array](): self.array().write_to(writer) - fn pretty_to[W: Writer](self, inout writer: W, indent: Variant[Int, String] = DefaultPrettyIndent): + fn pretty_to[W: Writer](self, inout writer: W, indent: String): if self.isa[Object](): - self.object().pretty_to(writer, indent) - if self.isa[Array](): - self.array().pretty_to(writer, indent) + writer.write("{\n") + self.object()._pretty_write_items(writer, indent * 2) + writer.write(indent, "}") + elif self.isa[Array](): + writer.write("[") + self.array()._pretty_write_items(writer, indent * 2) + writer.write(indent, "]") else: self.write_to(writer) diff --git a/test/ember_json/test_json.mojo b/test/ember_json/test_json.mojo index 3eb5244..5df7a6b 100644 --- a/test/ember_json/test_json.mojo +++ b/test/ember_json/test_json.mojo @@ -82,6 +82,18 @@ iamateapotnull ]""" assert_equal(expected, write_pretty(arr, indent=String("iamateapot"))) + arr = JSON.from_string('[123,"foo",false,{"key": null}]') + expected = """[ + 123, + "foo", + false, + { + "key": null + } +]""" + + assert_equal(expected, write_pretty(arr)) + def test_pretty_print_object(): var ob = JSON.from_string('{"k1": null, "k2": 123}') @@ -91,6 +103,17 @@ def test_pretty_print_object(): }""" assert_equal(expected, write_pretty(ob)) + ob = JSON.from_string('{"key": 123, "k": [123, false, null]}') + + expected = """{ + "key": 123, + "k": [ + 123, + false, + null + ] +}""" + def test_trailing_tokens(): with assert_raises(contains="Invalid json, expected end of input, recieved: garbage tokens"):