Add Encoder.CompactComments to omit extra new line (#541)

marshal.go

@@ -258,11 +258,12 @@ type Encoder struct {
 	w io.Writer
 	encOpts
 	annotation
-	line        int
-	col         int
-	order       MarshalOrder
-	promoteAnon bool
-	indentation string
+	line            int
+	col             int
+	order           MarshalOrder
+	promoteAnon     bool
+	compactComments bool
+	indentation     string
 }
 
 // NewEncoder returns a new encoder that writes to w.
@@ -369,6 +370,12 @@ func (e *Encoder) PromoteAnonymous(promote bool) *Encoder {
 	return e
 }
 
+// CompactComments removes the new line before each comment in the tree.
+func (e *Encoder) CompactComments(cc bool) *Encoder {
+	e.compactComments = cc
+	return e
+}
+
 func (e *Encoder) marshal(v interface{}) ([]byte, error) {
 	// Check if indentation is valid
 	for _, char := range e.indentation {
@@ -408,7 +415,7 @@ func (e *Encoder) marshal(v interface{}) ([]byte, error) {
 	}
 
 	var buf bytes.Buffer
-	_, err = t.writeToOrdered(&buf, "", "", 0, e.arraysOneElementPerLine, e.order, e.indentation, false)
+	_, err = t.writeToOrdered(&buf, "", "", 0, e.arraysOneElementPerLine, e.order, e.indentation, e.compactComments, false)
 
 	return buf.Bytes(), err
 }

marshal_test.go

@@ -1476,6 +1476,47 @@ commented out"""
 	}
 }
 
+func TestCompactComments(t *testing.T) {
+	expected := []byte(`
+[first-section]
+  # comment for first-key
+  first-key = 1
+  # comment for second-key
+  second-key = "value"
+  # comment for commented third-key
+  # third-key = []
+
+[second-section]
+  # comment for first-key
+  first-key = 2
+  # comment for second-key
+  second-key = "another value"
+  # comment for commented third-key
+  # third-key = ["value1", "value2"]
+`)
+	type Settings struct {
+		FirstKey  int      `toml:"first-key" comment:"comment for first-key"`
+		SecondKey string   `toml:"second-key" comment:"comment for second-key"`
+		ThirdKey  []string `toml:"third-key" comment:"comment for commented third-key" commented:"true"`
+	}
+	type Config struct {
+		FirstSection  Settings `toml:"first-section"`
+		SecondSection Settings `toml:"second-section"`
+	}
+	data := Config{
+		FirstSection:  Settings{1, "value", []string{}},
+		SecondSection: Settings{2, "another value", []string{"value1", "value2"}},
+	}
+	buf := new(bytes.Buffer)
+	if err := NewEncoder(buf).CompactComments(true).Encode(data); err != nil {
+		t.Fatal(err)
+	}
+
+	if !bytes.Equal(expected, buf.Bytes()) {
+		t.Errorf("Bad marshal: expected\n-----\n%s\n-----\ngot\n-----\n%s\n-----\n", expected, buf.Bytes())
+	}
+}
+
 type mapsTestStruct struct {
 	Simple map[string]string
 	Paths  map[string]string

tomltree_write.go

@@ -317,10 +317,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, false)
 }
 
-func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool, ord MarshalOrder, indentString string, parentCommented bool) (int64, error) {
+func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool, ord MarshalOrder, indentString string, compactComments, parentCommented bool) (int64, error) {
 	var orderedVals []sortNode
 
 	switch ord {
@@ -370,7 +370,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+indentString, combinedKey, bytesCount, arraysOneElementPerLine, ord, indentString, parentCommented || t.commented || tv.commented)
+				bytesCount, err = node.writeToOrdered(w, indent+indentString, combinedKey, bytesCount, arraysOneElementPerLine, ord, indentString, compactComments, parentCommented || t.commented || tv.commented)
 				if err != nil {
 					return bytesCount, err
 				}
@@ -386,7 +386,7 @@ func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount i
 						return bytesCount, err
 					}
 
-					bytesCount, err = subTree.writeToOrdered(w, indent+indentString, combinedKey, bytesCount, arraysOneElementPerLine, ord, indentString, parentCommented || t.commented || subTree.commented)
+					bytesCount, err = subTree.writeToOrdered(w, indent+indentString, combinedKey, bytesCount, arraysOneElementPerLine, ord, indentString, compactComments, parentCommented || t.commented || subTree.commented)
 					if err != nil {
 						return bytesCount, err
 					}
@@ -414,7 +414,14 @@ func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount i
 				if strings.HasPrefix(comment, "#") {
 					start = ""
 				}
-				writtenBytesCountComment, errc := writeStrings(w, "\n", indent, start, comment, "\n")
+				if !compactComments {
+					writtenBytesCountComment, errc := writeStrings(w, "\n")
+					bytesCount += int64(writtenBytesCountComment)
+					if errc != nil {
+						return bytesCount, errc
+					}
+				}
+				writtenBytesCountComment, errc := writeStrings(w, indent, start, comment, "\n")
 				bytesCount += int64(writtenBytesCountComment)
 				if errc != nil {
 					return bytesCount, errc