diff --git a/marshaler.go b/marshaler.go index 67513b0..b302893 100644 --- a/marshaler.go +++ b/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 diff --git a/marshaler_test.go b/marshaler_test.go index 8d64182..78d452e 100644 --- a/marshaler_test.go +++ b/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()