fix nested pretty print

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_as_element(writer, indent)
             if i < len(self._data) - 1:
                 writer.write(",")
             writer.write("\n")
 
-        writer.write("]")
-
     @always_inline
     fn __str__(self) -> String:
         return write(self)

ember_json/json.mojo

@@ -4,7 +4,7 @@ from .reader import Reader
 from .object import Object
 from .array import Array
 from .constants import *
-from .traits import JsonValue, PrettyPrintable, DefaultPrettyIndent
+from .traits import JsonValue, PrettyPrintable
 from .utils import write, ByteView
 
 
@@ -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:

ember_json/object.mojo

@@ -4,7 +4,7 @@ from collections import Dict
 from .constants import *
 from .utils import *
 from sys.intrinsics import unlikely, likely
-from .traits import JsonValue, PrettyPrintable, DefaultPrettyIndent
+from .traits import JsonValue, PrettyPrintable
 
 
 @value
@@ -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_as_element(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:

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):
         ...
 
 

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
 

ember_json/value.mojo

@@ -249,10 +249,22 @@ 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_as_element[W: Writer](self, inout writer: W, indent: String):
+        if self.isa[Object]():
+            writer.write("{\n")
+            self.object()._pretty_write_items(writer, indent * 2)
+            writer.write(indent, "}")
+        elif self.isa[Array]():
+            writer.write("[\n")
+            self.array()._pretty_write_items(writer, indent * 2)
+            writer.write(indent, "]")
+        else:
+            self.write_to(writer)
+
+    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]():
+        elif self.isa[Array]():
             self.array().pretty_to(writer, indent)
         else:
             self.write_to(writer)

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"):

test/ember_json/test_value.mojo

@@ -1,4 +1,5 @@
 from ember_json.value import Value, Null
+from ember_json.utils import write_pretty
 from testing import *
 
 def test_bool():
@@ -90,4 +91,21 @@ def test_bytes_for_string():
     assert_equal(Value("foobar").bytes_for_string(), 8)
     assert_equal(Value(Null()).bytes_for_string(), 4)
     assert_equal(Value(True).bytes_for_string(), 4)
-    assert_equal(Value(False).bytes_for_string(), 5)
+    assert_equal(Value(False).bytes_for_string(), 5)
+
+def test_pretty():
+    var v = Value.from_string("[123, 43564, false]")
+    var expected = """[
+    123,
+    43564,
+    false
+]"""
+    assert_equal(expected, write_pretty(v))
+
+    v = Value.from_string('{"key": 123, "k2": null}')
+    expected = """{
+    "key": 123,
+    "k2": null
+}"""
+
+    assert_equal(expected, write_pretty(v))