encoder: support indentation (#525)

marshaler.go

@@ -37,6 +37,7 @@ type Encoder struct {
 	tablesInline    bool
 	arraysMultiline bool
 	indentSymbol    string
+	indentTables    bool
 }
 
 // NewEncoder returns a new Encoder that writes to w.
@@ -64,6 +65,18 @@ func (enc *Encoder) SetArraysMultiline(multiline bool) {
 	enc.arraysMultiline = multiline
 }
 
+// SetIndentSymbol defines the string that should be used for indentation. The
+// provided string is repeated for each indentation level. Defaults to two
+// spaces.
+func (enc *Encoder) SetIndentSymbol(s string) {
+	enc.indentSymbol = s
+}
+
+// SetIndentTables forces the encoder to intent tables and array tables.
+func (enc *Encoder) SetIndentTables(indent bool) {
+	enc.indentTables = indent
+}
+
 // Encode writes a TOML representation of v to the stream.
 //
 // If v cannot be represented to TOML it returns an error.
@@ -257,6 +270,8 @@ func (enc *Encoder) encodeKv(b []byte, ctx encoderCtx, options valueOptions, v r
 		return b, nil
 	}
 
+	b = enc.indent(ctx.indent, b)
+
 	b, err = enc.encodeKey(b, ctx.key)
 	if err != nil {
 		return nil, err
@@ -367,21 +382,23 @@ func (enc *Encoder) encodeUnquotedKey(b []byte, v string) []byte {
 	return append(b, v...)
 }
 
-func (enc *Encoder) encodeTableHeader(b []byte, key []string) ([]byte, error) {
-	if len(key) == 0 {
+func (enc *Encoder) encodeTableHeader(ctx encoderCtx, b []byte) ([]byte, error) {
+	if len(ctx.parentKey) == 0 {
 		return b, nil
 	}
 
+	b = enc.indent(ctx.indent, b)
+
 	b = append(b, '[')
 
 	var err error
 
-	b, err = enc.encodeKey(b, key[0])
+	b, err = enc.encodeKey(b, ctx.parentKey[0])
 	if err != nil {
 		return nil, err
 	}
 
-	for _, k := range key[1:] {
+	for _, k := range ctx.parentKey[1:] {
 		b = append(b, '.')
 
 		b, err = enc.encodeKey(b, k)
@@ -559,10 +576,13 @@ func (enc *Encoder) encodeTable(b []byte, ctx encoderCtx, t table) ([]byte, erro
 	}
 
 	if !ctx.skipTableHeader {
-		b, err = enc.encodeTableHeader(b, ctx.parentKey)
+		b, err = enc.encodeTableHeader(ctx, b)
 		if err != nil {
 			return nil, err
 		}
+		if enc.indentTables && len(ctx.parentKey) > 0 {
+			ctx.indent++
+		}
 	}
 	ctx.skipTableHeader = false
 

marshaler_test.go

@@ -348,6 +348,7 @@ type flagsSetters []struct {
 var allFlags = flagsSetters{
 	{"arrays-multiline", (*toml.Encoder).SetArraysMultiline},
 	{"tables-inline", (*toml.Encoder).SetTablesInline},
+	{"indent-tables", (*toml.Encoder).SetIndentTables},
 }
 
 func setFlags(enc *toml.Encoder, flags int) {
@@ -393,6 +394,66 @@ func equalStringsIgnoreNewlines(t *testing.T, expected string, actual string) {
 	assert.Equal(t, strings.Trim(expected, cutset), strings.Trim(actual, cutset))
 }
 
+func TestMarshalIndentTables(t *testing.T) {
+	examples := []struct {
+		desc     string
+		v        interface{}
+		expected string
+	}{
+		{
+			desc: "one kv",
+			v: map[string]interface{}{
+				"foo": "bar",
+			},
+			expected: `foo = 'bar'`,
+		},
+		{
+			desc: "one level table",
+			v: map[string]map[string]string{
+				"foo": {
+					"one": "value1",
+					"two": "value2",
+				},
+			},
+			expected: `
+[foo]
+  one = 'value1'
+  two = 'value2'
+`,
+		},
+		{
+			desc: "two levels table",
+			v: map[string]interface{}{
+				"root": "value0",
+				"level1": map[string]interface{}{
+					"one": "value1",
+					"level2": map[string]interface{}{
+						"two": "value2",
+					},
+				},
+			},
+			expected: `
+root = 'value0'
+[level1]
+  one = 'value1'
+  [level1.level2]
+    two = 'value2'
+`,
+		},
+	}
+
+	for _, e := range examples {
+		t.Run(e.desc, func(t *testing.T) {
+			var buf strings.Builder
+			enc := toml.NewEncoder(&buf)
+			enc.SetIndentTables(true)
+			err := enc.Encode(e.v)
+			require.NoError(t, err)
+			equalStringsIgnoreNewlines(t, e.expected, buf.String())
+		})
+	}
+}
+
 func TestIssue436(t *testing.T) {
 	t.Parallel()