From 82a6a1977d0361be1013756fae1c0a687eec4fd4 Mon Sep 17 00:00:00 2001
From: x-hgg-x <39058530+x-hgg-x@users.noreply.github.com>
Date: Mon, 4 May 2020 19:21:22 +0200
Subject: [PATCH] Add indentation setting for Encoder (#386)

Fixes #371
---
 marshal.go        | 29 ++++++++++++++++++++++-------
 marshal_test.go   | 33 +++++++++++++++++++++++++++++++++
 tomltree_write.go |  8 ++++----
 3 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/marshal.go b/marshal.go
index a9ee1c1..6917ea2 100644
--- a/marshal.go
+++ b/marshal.go
@@ -218,17 +218,19 @@ type Encoder struct {
 	col         int
 	order       marshalOrder
 	promoteAnon bool
+	indentation string
 }
 
 // NewEncoder returns a new encoder that writes to w.
 func NewEncoder(w io.Writer) *Encoder {
 	return &Encoder{
-		w:          w,
-		encOpts:    encOptsDefaults,
-		annotation: annotationDefault,
-		line:       0,
-		col:        1,
-		order:      OrderAlphabetical,
+		w:           w,
+		encOpts:     encOptsDefaults,
+		annotation:  annotationDefault,
+		line:        0,
+		col:         1,
+		order:       OrderAlphabetical,
+		indentation: "  ",
 	}
 }
 
@@ -280,6 +282,12 @@ func (e *Encoder) Order(ord marshalOrder) *Encoder {
 	return e
 }
 
+// Indentation allows to change indentation when marshalling.
+func (e *Encoder) Indentation(indent string) *Encoder {
+	e.indentation = indent
+	return e
+}
+
 // SetTagName allows changing default tag "toml"
 func (e *Encoder) SetTagName(v string) *Encoder {
 	e.tag = v
@@ -318,6 +326,13 @@ func (e *Encoder) PromoteAnonymous(promote bool) *Encoder {
 }
 
 func (e *Encoder) marshal(v interface{}) ([]byte, error) {
+	// Check if indentation is valid
+	for _, char := range e.indentation {
+		if !(char == ' ' || char == '\t') {
+			return []byte{}, fmt.Errorf("invalid indentation: must only contains space or tab characters")
+		}
+	}
+
 	mtype := reflect.TypeOf(v)
 	if mtype == nil {
 		return []byte{}, errors.New("nil cannot be marshaled to TOML")
@@ -349,7 +364,7 @@ func (e *Encoder) marshal(v interface{}) ([]byte, error) {
 	}
 
 	var buf bytes.Buffer
-	_, err = t.writeToOrdered(&buf, "", "", 0, e.arraysOneElementPerLine, e.order, false)
+	_, err = t.writeToOrdered(&buf, "", "", 0, e.arraysOneElementPerLine, e.order, e.indentation, false)
 
 	return buf.Bytes(), err
 }
diff --git a/marshal_test.go b/marshal_test.go
index 7d6c025..5ef49a5 100644
--- a/marshal_test.go
+++ b/marshal_test.go
@@ -44,6 +44,19 @@ Zstring = "Hello"
   String2 = "One"
 `)
 
+var basicTestTomlCustomIndentation = []byte(`Ystrlist = ["Howdy","Hey There"]
+Zstring = "Hello"
+
+[[Wsublist]]
+	String2 = "Two"
+
+[[Wsublist]]
+	String2 = "Three"
+
+[Xsubdoc]
+	String2 = "One"
+`)
+
 var basicTestTomlOrdered = []byte(`Zstring = "Hello"
 Ystrlist = ["Howdy","Hey There"]
 
@@ -207,6 +220,26 @@ func TestBasicMarshal(t *testing.T) {
 	}
 }
 
+func TestBasicMarshalCustomIndentation(t *testing.T) {
+	var result bytes.Buffer
+	err := NewEncoder(&result).Indentation("\t").Encode(basicTestData)
+	if err != nil {
+		t.Fatal(err)
+	}
+	expected := basicTestTomlCustomIndentation
+	if !bytes.Equal(result.Bytes(), expected) {
+		t.Errorf("Bad marshal: expected\n-----\n%s\n-----\ngot\n-----\n%s\n-----\n", expected, result.Bytes())
+	}
+}
+
+func TestBasicMarshalWrongIndentation(t *testing.T) {
+	var result bytes.Buffer
+	err := NewEncoder(&result).Indentation("  \n").Encode(basicTestData)
+	if err.Error() != "invalid indentation: must only contains space or tab characters" {
+		t.Error("expect err:invalid indentation: must only contains space or tab characters but got:", err)
+	}
+}
+
 func TestBasicMarshalOrdered(t *testing.T) {
 	var result bytes.Buffer
 	err := NewEncoder(&result).Order(OrderPreserve).Encode(basicTestData)
diff --git a/tomltree_write.go b/tomltree_write.go
index 9acc2f3..5d2faff 100644
--- a/tomltree_write.go
+++ b/tomltree_write.go
@@ -282,10 +282,10 @@ func sortAlphabetical(t *Tree) (vals []sortNode) {
 }
 
 func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool) (int64, error) {
-	return t.writeToOrdered(w, indent, keyspace, bytesCount, arraysOneElementPerLine, OrderAlphabetical, false)
+	return t.writeToOrdered(w, indent, keyspace, bytesCount, arraysOneElementPerLine, OrderAlphabetical, "  ", false)
 }
 
-func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool, ord marshalOrder, parentCommented bool) (int64, error) {
+func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool, ord marshalOrder, indentString string, parentCommented bool) (int64, error) {
 	var orderedVals []sortNode
 
 	switch ord {
@@ -335,7 +335,7 @@ func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount i
 				if err != nil {
 					return bytesCount, err
 				}
-				bytesCount, err = node.writeToOrdered(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine, ord, parentCommented || t.commented || tv.commented)
+				bytesCount, err = node.writeToOrdered(w, indent+indentString, combinedKey, bytesCount, arraysOneElementPerLine, ord, indentString, parentCommented || t.commented || tv.commented)
 				if err != nil {
 					return bytesCount, err
 				}
@@ -351,7 +351,7 @@ func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount i
 						return bytesCount, err
 					}
 
-					bytesCount, err = subTree.writeToOrdered(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine, ord, parentCommented || t.commented || subTree.commented)
+					bytesCount, err = subTree.writeToOrdered(w, indent+indentString, combinedKey, bytesCount, arraysOneElementPerLine, ord, indentString, parentCommented || t.commented || subTree.commented)
 					if err != nil {
 						return bytesCount, err
 					}