From d08347058532c70f4c6b6821bf8d3662d3f19bbf Mon Sep 17 00:00:00 2001 From: "Mikhail f. Shiryaev" Date: Wed, 12 May 2021 15:22:40 +0200 Subject: [PATCH] Add Encoder.CompactComments to omit extra new line (#541) --- marshal.go | 19 +++++++++++++------ marshal_test.go | 41 +++++++++++++++++++++++++++++++++++++++++ tomltree_write.go | 17 ++++++++++++----- 3 files changed, 66 insertions(+), 11 deletions(-) diff --git a/marshal.go b/marshal.go index 53f4d4e2..3443c354 100644 --- a/marshal.go +++ b/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 } diff --git a/marshal_test.go b/marshal_test.go index 48f7df76..1abd5c1a 100644 --- a/marshal_test.go +++ b/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 diff --git a/tomltree_write.go b/tomltree_write.go index 4e78c0b5..c9afbdab 100644 --- a/tomltree_write.go +++ b/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