From 9ba52996d8336c93a92b8021ae6ee3a3bbbba142 Mon Sep 17 00:00:00 2001
From: Thomas Pelletier <pelletier.thomas@gmail.com>
Date: Wed, 21 Apr 2021 22:13:45 -0400
Subject: [PATCH] Encoder multiline array (#520)

---
 marshaler.go      | 59 ++++++++++++++++++++++++++++----
 marshaler_test.go | 85 ++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 133 insertions(+), 11 deletions(-)

diff --git a/marshaler.go b/marshaler.go
index 3153108..2491015 100644
--- a/marshaler.go
+++ b/marshaler.go
@@ -33,21 +33,36 @@ type Encoder struct {
 	w io.Writer
 
 	// global settings
-	tablesInline bool
+	tablesInline    bool
+	arraysMultiline bool
+	indentSymbol    string
 }
 
 // NewEncoder returns a new Encoder that writes to w.
 func NewEncoder(w io.Writer) *Encoder {
 	return &Encoder{
-		w: w,
+		w:            w,
+		indentSymbol: "  ",
 	}
 }
 
 // SetTablesInline forces the encoder to emit all tables inline.
+//
+// This behavior can be controled on an individual struct field basis with the
+// `inline="true"` tag.
 func (e *Encoder) SetTablesInline(inline bool) {
 	e.tablesInline = inline
 }
 
+// SetArraysMultiline forces the encoder to emit all arrays with one element per
+// line.
+//
+// This behavior can be controled on an individual struct field basis with the
+// `multiline="true"` tag.
+func (e *Encoder) SetArraysMultiline(multiline bool) {
+	e.arraysMultiline = multiline
+}
+
 // Encode writes a TOML representation of v to the stream.
 //
 // If v cannot be represented to TOML it returns an error.
@@ -121,6 +136,10 @@ type encoderCtx struct {
 	// Should the next table be encoded as inline
 	inline bool
 
+	// Indentation level
+	indent int
+
+	// Options coming from struct tags
 	options valueOptions
 }
 
@@ -544,6 +563,7 @@ func (enc *Encoder) encodeTable(b []byte, ctx encoderCtx, t table) ([]byte, erro
 	for _, table := range t.tables {
 		ctx.setKey(table.Key)
 
+		ctx.options = table.Options
 		b, err = enc.encode(b, ctx, table.Value)
 		if err != nil {
 			return nil, err
@@ -721,25 +741,52 @@ func (enc *Encoder) encodeSliceAsArrayTable(b []byte, ctx encoderCtx, v reflect.
 }
 
 func (enc *Encoder) encodeSliceAsArray(b []byte, ctx encoderCtx, v reflect.Value) ([]byte, error) {
+	multiline := ctx.options.multiline || enc.arraysMultiline
+	separator := ", "
+
 	b = append(b, '[')
 
+	subCtx := ctx
+	subCtx.options = valueOptions{}
+
+	if multiline {
+		separator = ",\n"
+		b = append(b, '\n')
+		subCtx.indent++
+	}
+
 	var err error
 	first := true
 
 	for i := 0; i < v.Len(); i++ {
-		if !first {
-			b = append(b, ", "...)
+		if first {
+			first = false
+		} else {
+			b = append(b, separator...)
 		}
 
-		first = false
+		if multiline {
+			b = enc.indent(subCtx.indent, b)
+		}
 
-		b, err = enc.encode(b, ctx, v.Index(i))
+		b, err = enc.encode(b, subCtx, v.Index(i))
 		if err != nil {
 			return nil, err
 		}
 	}
 
+	if multiline {
+		b = append(b, '\n')
+		b = enc.indent(ctx.indent, b)
+	}
 	b = append(b, ']')
 
 	return b, nil
 }
+
+func (enc *Encoder) indent(level int, b []byte) []byte {
+	for i := 0; i < level; i++ {
+		b = append(b, enc.indentSymbol...)
+	}
+	return b
+}
diff --git a/marshaler_test.go b/marshaler_test.go
index 3304f17..646de24 100644
--- a/marshaler_test.go
+++ b/marshaler_test.go
@@ -3,6 +3,7 @@ package toml_test
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"strings"
 	"testing"
 
@@ -262,6 +263,39 @@ world"""`,
 A = {isinline = 'yes'}
 [B]
 isinline = 'no'
+`,
+		},
+		{
+			desc: "mutiline array int",
+			v: struct {
+				A []int `multiline:"true"`
+				B []int
+			}{
+				A: []int{1, 2, 3, 4},
+				B: []int{1, 2, 3, 4},
+			},
+			expected: `
+A = [
+  1,
+  2,
+  3,
+  4
+]
+B = [1, 2, 3, 4]
+`,
+		},
+		{
+			desc: "mutiline array in array",
+			v: struct {
+				A [][]int `multiline:"true"`
+			}{
+				A: [][]int{{1, 2}, {3, 4}},
+			},
+			expected: `
+A = [
+  [1, 2],
+  [3, 4]
+]
 `,
 		},
 	}
@@ -285,13 +319,10 @@ isinline = 'no'
 			err = toml.Unmarshal(b, &defaultMap)
 			require.NoError(t, err)
 
-			// checks that the TablesInline mode generates valid,
-			// equivalent TOML
-			t.Run("tables inline", func(t *testing.T) {
+			testWithAllFlags(t, func(t *testing.T, flags int) {
 				var buf bytes.Buffer
-
 				enc := toml.NewEncoder(&buf)
-				enc.SetTablesInline(true)
+				setFlags(enc, flags)
 
 				err := enc.Encode(e.v)
 				require.NoError(t, err)
@@ -306,6 +337,50 @@ isinline = 'no'
 	}
 }
 
+type flagsSetters []struct {
+	name string
+	f    func(enc *toml.Encoder, flag bool)
+}
+
+var allFlags = flagsSetters{
+	{"arrays-multiline", (*toml.Encoder).SetArraysMultiline},
+	{"tables-inline", (*toml.Encoder).SetTablesInline},
+}
+
+func setFlags(enc *toml.Encoder, flags int) {
+	for i := 0; i < len(allFlags); i++ {
+		enabled := flags&1 > 0
+		allFlags[i].f(enc, enabled)
+	}
+}
+
+func testWithAllFlags(t *testing.T, testfn func(t *testing.T, flags int)) {
+	t.Helper()
+	testWithFlags(t, 0, allFlags, testfn)
+}
+
+func testWithFlags(t *testing.T, flags int, setters flagsSetters, testfn func(t *testing.T, flags int)) {
+	t.Helper()
+
+	if len(setters) == 0 {
+		testfn(t, flags)
+		return
+	}
+
+	s := setters[0]
+
+	for _, enabled := range []bool{false, true} {
+		name := fmt.Sprintf("%s=%t", s.name, enabled)
+		newFlags := flags << 1
+		if enabled {
+			newFlags++
+		}
+		t.Run(name, func(t *testing.T) {
+			testWithFlags(t, newFlags, setters[1:], testfn)
+		})
+	}
+}
+
 func equalStringsIgnoreNewlines(t *testing.T, expected string, actual string) {
 	t.Helper()
 	cutset := "\n"