Repository: amzn/ion-go
Pull request #32 (2020-05-20 12:32:50)

Adding support for pretty-printing text Ion
By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.




Patch:
diff --git a/ion/textwriter.go b/ion/textwriter.go
index 2ae82c1..282ca8f 100644
--- a/ion/textwriter.go
+++ b/ion/textwriter.go
@@ -16,13 +16,19 @@ const (
 	// know you're only emiting one datagram; dangerous if there's a chance you're going
 	// to emit another datagram using the same Writer.
 	TextWriterQuietFinish TextWriterOpts = 1
+
+	// TextWriterPretty enables pretty-printing mode.
+	TextWriterPretty TextWriterOpts = 2
 )
 
 // textWriter is a writer that writes human-readable text
 type textWriter struct {
 	writer
-	needsSeparator bool
 	opts           TextWriterOpts
+	needsSeparator bool
+	emptyContainer bool
+	emptyStream    bool
+	indent         int
 }
 
 // NewTextWriter returns a new text writer.
@@ -33,10 +39,9 @@ func NewTextWriter(out io.Writer) Writer {
 // NewTextWriterOpts returns a new text writer with the given options.
 func NewTextWriterOpts(out io.Writer, opts TextWriterOpts) Writer {
 	return &textWriter{
-		writer: writer{
-			out: out,
-		},
-		opts: opts,
+		writer:      writer{out: out},
+		opts:        opts,
+		emptyStream: true,
 	}
 }
 
@@ -244,17 +249,23 @@ func (w *textWriter) Finish() error {
 		return &UsageError{"Writer.Finish", "not at top level"}
 	}
 
-	if w.opts&TextWriterQuietFinish == 0 {
+	if !w.emptyStream && w.opts&TextWriterQuietFinish == 0 {
 		if w.err = writeRawChar('\n', w.out); w.err != nil {
 			return w.err
 		}
 		w.needsSeparator = false
+		w.emptyStream = true
 	}
 
 	w.clear()
 	return nil
 }
 
+// pretty returns true if we're pretty-printing.
+func (w *textWriter) pretty() bool {
+	return w.opts&TextWriterPretty == TextWriterPretty
+}
+
 // writeValue writes a stringified value to the output stream.
 func (w *textWriter) writeValue(api string, val string) error {
 	if w.err != nil {
@@ -277,47 +288,100 @@ func (w *textWriter) writeValue(api string, val string) error {
 // annotations (if any).
 func (w *textWriter) beginValue(api string) error {
 	if w.needsSeparator {
-		var sep byte
-		switch w.ctx.peek() {
-		case ctxInStruct, ctxInList:
-			sep = ','
-		case ctxInSexp:
-			sep = ' '
-		default:
-			sep = '\n'
+		if err := w.writeSeparator(); err != nil {
+			return err
 		}
+	}
 
-		if err := writeRawChar(sep, w.out); err != nil {
+	if w.emptyContainer {
+		if w.pretty() {
+			if err := writeRawChar('\n', w.out); err != nil {
+				return err
+			}
+		}
+	}
+
+	if w.pretty() {
+		if err := w.writeIndent(); err != nil {
 			return err
 		}
 	}
 
 	if w.inStruct() {
-		if w.fieldName == "" {
-			return &UsageError{api, "field name not set"}
+		if err := w.writeFieldName(api); err != nil {
+			return err
 		}
-		name := w.fieldName
-		w.fieldName = ""
+	}
 
-		if err := writeSymbol(name, w.out); err != nil {
+	if len(w.annotations) > 0 {
+		if err := w.writeAnnotations(); err != nil {
 			return err
 		}
-		if err := writeRawChar(':', w.out); err != nil {
-			return err
+	}
+
+	return nil
+}
+
+// writeSeparator writes out the character or characters that separate values.
+func (w *textWriter) writeSeparator() error {
+	var sep string
+
+	switch w.ctx.peek() {
+	case ctxInStruct, ctxInList:
+		// In a struct or a list, values are separated by commas.
+		if w.pretty() {
+			sep = ",\n"
+		} else {
+			sep = ","
 		}
+
+	case ctxInSexp:
+		// In an sexp, values are separated by whitespace.
+		if w.pretty() {
+			sep = "\n"
+		} else {
+			sep = " "
+		}
+
+	default:
+		// At the top level, values are separated by newlines.
+		sep = "\n"
 	}
 
-	if len(w.annotations) > 0 {
-		as := w.annotations
-		w.annotations = nil
+	return writeRawString(sep, w.out)
+}
 
-		for _, a := range as {
-			if err := writeSymbol(a, w.out); err != nil {
-				return err
-			}
-			if err := writeRawString("::", w.out); err != nil {
-				return err
-			}
+// writeFieldName writes a field name inside a struct.
+func (w *textWriter) writeFieldName(api string) error {
+	if w.fieldName == "" {
+		return &UsageError{api, "field name not set"}
+	}
+	name := w.fieldName
+	w.fieldName = ""
+
+	if err := writeSymbol(name, w.out); err != nil {
+		return err
+	}
+
+	sep := ":"
+	if w.pretty() {
+		sep = ": "
+	}
+
+	return writeRawString(sep, w.out)
+}
+
+// writeAnnotations writes out the annotations for a value.
+func (w *textWriter) writeAnnotations() error {
+	as := w.annotations
+	w.annotations = nil
+
+	for _, a := range as {
+		if err := writeSymbol(a, w.out); err != nil {
+			return err
+		}
+		if err := writeRawString("::", w.out); err != nil {
+			return err
 		}
 	}
 
@@ -327,6 +391,8 @@ func (w *textWriter) beginValue(api string) error {
 // endValue finishes the process of writing a value.
 func (w *textWriter) endValue() {
 	w.needsSeparator = true
+	w.emptyContainer = false
+	w.emptyStream = false
 }
 
 // begin starts writing a container of the given type.
@@ -336,7 +402,9 @@ func (w *textWriter) begin(api string, t ctx, c byte) error {
 	}
 
 	w.ctx.push(t)
+	w.indent++
 	w.needsSeparator = false
+	w.emptyContainer = true
 
 	return writeRawChar(c, w.out)
 }
@@ -347,6 +415,17 @@ func (w *textWriter) end(api string, t ctx, c byte) error {
 		return &UsageError{api, "not in that kind of container"}
 	}
 
+	w.indent--
+
+	if !w.emptyContainer && w.pretty() {
+		if err := writeRawChar('\n', w.out); err != nil {
+			return err
+		}
+		if err := w.writeIndent(); err != nil {
+			return err
+		}
+	}
+
 	if err := writeRawChar(c, w.out); err != nil {
 		return err
 	}
@@ -357,3 +436,13 @@ func (w *textWriter) end(api string, t ctx, c byte) error {
 
 	return nil
 }
+
+// writeIndent writes out tabs to indent a pretty-printed value.
+func (w *textWriter) writeIndent() error {
+	for i := 0; i < w.indent; i++ {
+		if err := writeRawChar('\t', w.out); err != nil {
+			return err
+		}
+	}
+	return nil
+}


Test patch:
diff --git a/ion/textwriter_test.go b/ion/textwriter_test.go
index b3124c3..2aa7671 100644
--- a/ion/textwriter_test.go
+++ b/ion/textwriter_test.go
@@ -334,6 +334,81 @@ func TestWriteTextBadFinish(t *testing.T) {
 	}
 }
 
+func TestWriteTextPretty(t *testing.T) {
+	buf := strings.Builder{}
+	w := NewTextWriterOpts(&buf, TextWriterPretty)
+
+	w.BeginStruct()
+	{
+		w.FieldName("struct")
+		w.BeginStruct()
+		w.EndStruct()
+
+		w.FieldName("list")
+		w.Annotations("i", "am", "a", "list")
+		w.BeginList()
+		{
+			w.WriteString("value")
+			w.WriteNullType(StringType)
+			w.BeginStruct()
+			{
+				w.FieldName("1")
+				w.WriteString("one")
+				w.FieldName("2")
+				w.WriteString("two")
+			}
+			w.EndStruct()
+		}
+		w.EndList()
+
+		w.FieldName("sexp")
+		w.BeginSexp()
+		{
+			w.WriteSymbol("+")
+			w.WriteInt(123)
+			w.BeginSexp()
+			{
+				w.WriteSymbol("*")
+				w.WriteInt(456)
+				w.WriteInt(789)
+			}
+			w.EndSexp()
+		}
+		w.EndSexp()
+	}
+	w.EndStruct()
+
+	if err := w.Finish(); err != nil {
+		t.Fatal(err)
+	}
+
+	actual := buf.String()
+	expected := `{
+	struct: {},
+	list: i::am::a::list::[
+		"value",
+		null.string,
+		{
+			'1': "one",
+			'2': "two"
+		}
+	],
+	sexp: (
+		'+'
+		123
+		(
+			'*'
+			456
+			789
+		)
+	)
+}
+`
+	if actual != expected {
+		t.Errorf("expected:\n%v\ngot:\n%v", expected, actual)
+	}
+}
+
 func testTextWriter(t *testing.T, expected string, f func(Writer)) {
 	actual := writeText(f)
 	if actual != expected {