Fix multiline + non-primitive commenting (#336)

Fixes #216

marshal.go

@@ -302,7 +302,7 @@ func (e *Encoder) marshal(v interface{}) ([]byte, error) {
 	}
 
 	var buf bytes.Buffer
-	_, err = t.writeToOrdered(&buf, "", "", 0, e.arraysOneElementPerLine, e.order)
+	_, err = t.writeToOrdered(&buf, "", "", 0, e.arraysOneElementPerLine, e.order, false)
 
 	return buf.Bytes(), err
 }
@@ -363,7 +363,7 @@ func (e *Encoder) valueToTree(mtype reflect.Type, mval reflect.Value) (*Tree, er
 				return nil, err
 			}
 			if e.quoteMapKeys {
-				keyStr, err := tomlValueStringRepresentation(key.String(), "", e.arraysOneElementPerLine)
+				keyStr, err := tomlValueStringRepresentation(key.String(), "", "", e.arraysOneElementPerLine)
 				if err != nil {
 					return nil, err
 				}

marshal_test.go

@@ -951,6 +951,213 @@ func TestMarshalComment(t *testing.T) {
 	}
 }
 
+func TestMarshalMultilineCommented(t *testing.T) {
+	expectedToml := []byte(`# MultilineArray = [
+  # 100,
+  # 200,
+  # 300,
+# ]
+# MultilineNestedArray = [
+  # [
+  # "a",
+  # "b",
+  # "c",
+# ],
+  # [
+  # "d",
+  # "e",
+  # "f",
+# ],
+# ]
+# MultilineString = """
+# I
+# am
+# Allen"""
+NonCommented = "Not commented line"
+`)
+	type StructWithMultiline struct {
+		NonCommented         string
+		MultilineString      string     `commented:"true" multiline:"true"`
+		MultilineArray       []int      `commented:"true"`
+		MultilineNestedArray [][]string `commented:"true"`
+	}
+
+	var buf bytes.Buffer
+	enc := NewEncoder(&buf)
+	if err := enc.ArraysWithOneElementPerLine(true).Encode(StructWithMultiline{
+		NonCommented:    "Not commented line",
+		MultilineString: "I\nam\nAllen",
+		MultilineArray:  []int{100, 200, 300},
+		MultilineNestedArray: [][]string{
+			{"a", "b", "c"},
+			{"d", "e", "f"},
+		},
+	}); err == nil {
+		result := buf.Bytes()
+		if !bytes.Equal(result, expectedToml) {
+			t.Errorf("Bad marshal: expected\n-----\n%s\n-----\ngot\n-----\n%s\n-----\n", expectedToml, result)
+		}
+	} else {
+		t.Fatal(err)
+	}
+}
+
+func TestMarshalNonPrimitiveTypeCommented(t *testing.T) {
+	expectedToml := []byte(`
+# [CommentedMapField]
+
+  # [CommentedMapField.CommentedMapField1]
+    # SingleLineString = "This line should be commented out"
+
+  # [CommentedMapField.CommentedMapField2]
+    # SingleLineString = "This line should be commented out"
+
+# [CommentedStructField]
+
+  # [CommentedStructField.CommentedStructField]
+    # MultilineArray = [
+      # 1,
+      # 2,
+    # ]
+    # MultilineNestedArray = [
+      # [
+      # 10,
+      # 20,
+    # ],
+      # [
+      # 100,
+      # 200,
+    # ],
+    # ]
+    # MultilineString = """
+# This line
+# should be
+# commented out"""
+
+  # [CommentedStructField.NotCommentedStructField]
+    # MultilineArray = [
+      # 1,
+      # 2,
+    # ]
+    # MultilineNestedArray = [
+      # [
+      # 10,
+      # 20,
+    # ],
+      # [
+      # 100,
+      # 200,
+    # ],
+    # ]
+    # MultilineString = """
+# This line
+# should be
+# commented out"""
+
+[NotCommentedStructField]
+
+  # [NotCommentedStructField.CommentedStructField]
+    # MultilineArray = [
+      # 1,
+      # 2,
+    # ]
+    # MultilineNestedArray = [
+      # [
+      # 10,
+      # 20,
+    # ],
+      # [
+      # 100,
+      # 200,
+    # ],
+    # ]
+    # MultilineString = """
+# This line
+# should be
+# commented out"""
+
+  [NotCommentedStructField.NotCommentedStructField]
+    MultilineArray = [
+      3,
+      4,
+    ]
+    MultilineNestedArray = [
+      [
+      30,
+      40,
+    ],
+      [
+      300,
+      400,
+    ],
+    ]
+    MultilineString = """
+This line
+should NOT be
+commented out"""
+`)
+	type InnerStruct struct {
+		MultilineString      string `multiline:"true"`
+		MultilineArray       []int
+		MultilineNestedArray [][]int
+	}
+	type MiddleStruct struct {
+		NotCommentedStructField InnerStruct
+		CommentedStructField    InnerStruct `commented:"true"`
+	}
+	type OuterStruct struct {
+		CommentedStructField    MiddleStruct `commented:"true"`
+		NotCommentedStructField MiddleStruct
+		CommentedMapField       map[string]struct{ SingleLineString string } `commented:"true"`
+	}
+
+	commentedTestStruct := OuterStruct{
+		CommentedStructField: MiddleStruct{
+			NotCommentedStructField: InnerStruct{
+				MultilineString:      "This line\nshould be\ncommented out",
+				MultilineArray:       []int{1, 2},
+				MultilineNestedArray: [][]int{{10, 20}, {100, 200}},
+			},
+			CommentedStructField: InnerStruct{
+				MultilineString:      "This line\nshould be\ncommented out",
+				MultilineArray:       []int{1, 2},
+				MultilineNestedArray: [][]int{{10, 20}, {100, 200}},
+			},
+		},
+		NotCommentedStructField: MiddleStruct{
+			NotCommentedStructField: InnerStruct{
+				MultilineString:      "This line\nshould NOT be\ncommented out",
+				MultilineArray:       []int{3, 4},
+				MultilineNestedArray: [][]int{{30, 40}, {300, 400}},
+			},
+			CommentedStructField: InnerStruct{
+				MultilineString:      "This line\nshould be\ncommented out",
+				MultilineArray:       []int{1, 2},
+				MultilineNestedArray: [][]int{{10, 20}, {100, 200}},
+			},
+		},
+		CommentedMapField: map[string]struct{ SingleLineString string }{
+			"CommentedMapField1": {
+				SingleLineString: "This line should be commented out",
+			},
+			"CommentedMapField2": {
+				SingleLineString: "This line should be commented out",
+			},
+		},
+	}
+
+	var buf bytes.Buffer
+	enc := NewEncoder(&buf)
+	if err := enc.ArraysWithOneElementPerLine(true).Encode(commentedTestStruct); err == nil {
+		result := buf.Bytes()
+		if !bytes.Equal(result, expectedToml) {
+			t.Errorf("Bad marshal: expected\n-----\n%s\n-----\ngot\n-----\n%s\n-----\n", expectedToml, result)
+		}
+	} else {
+		t.Fatal(err)
+	}
+}
+
 type mapsTestStruct struct {
 	Simple map[string]string
 	Paths  map[string]string

parser_test.go

@@ -897,7 +897,7 @@ func TestTomlValueStringRepresentation(t *testing.T) {
 			"[\"gamma\",\"delta\"]"},
 		{nil, ""},
 	} {
-		result, err := tomlValueStringRepresentation(item.Value, "", false)
+		result, err := tomlValueStringRepresentation(item.Value, "", "", false)
 		if err != nil {
 			t.Errorf("Test %d - unexpected error: %s", idx, err)
 		}

toml.go

@@ -222,8 +222,12 @@ func (t *Tree) SetPathWithOptions(keys []string, opts SetOptions, value interfac
 	switch v := value.(type) {
 	case *Tree:
 		v.comment = opts.Comment
+		v.commented = opts.Commented
 		toInsert = value
 	case []*Tree:
+		for i := range v {
+			v[i].commented = opts.Commented
+		}
 		toInsert = value
 	case *tomlValue:
 		v.comment = opts.Comment

tomltree_write.go

@@ -28,9 +28,10 @@ type sortNode struct {
 // Encodes a string to a TOML-compliant multi-line string value
 // This function is a clone of the existing encodeTomlString function, except that whitespace characters
 // are preserved. Quotation marks and backslashes are also not escaped.
-func encodeMultilineTomlString(value string) string {
+func encodeMultilineTomlString(value string, commented string) string {
 	var b bytes.Buffer
 
+	b.WriteString(commented)
 	for _, rr := range value {
 		switch rr {
 		case '\b':
@@ -38,7 +39,7 @@ func encodeMultilineTomlString(value string) string {
 		case '\t':
 			b.WriteString("\t")
 		case '\n':
-			b.WriteString("\n")
+			b.WriteString("\n" + commented)
 		case '\f':
 			b.WriteString(`\f`)
 		case '\r':
@@ -91,7 +92,7 @@ func encodeTomlString(value string) string {
 	return b.String()
 }
 
-func tomlValueStringRepresentation(v interface{}, indent string, arraysOneElementPerLine bool) (string, error) {
+func tomlValueStringRepresentation(v interface{}, commented string, indent string, arraysOneElementPerLine bool) (string, error) {
 	// this interface check is added to dereference the change made in the writeTo function.
 	// That change was made to allow this function to see formatting options.
 	tv, ok := v.(*tomlValue)
@@ -123,12 +124,12 @@ func tomlValueStringRepresentation(v interface{}, indent string, arraysOneElemen
 		return strings.ToLower(strconv.FormatFloat(value, 'f', -1, bits)), nil
 	case string:
 		if tv.multiline {
-			return "\"\"\"\n" + encodeMultilineTomlString(value) + "\"\"\"", nil
+			return "\"\"\"\n" + encodeMultilineTomlString(value, commented) + "\"\"\"", nil
 		}
 		return "\"" + encodeTomlString(value) + "\"", nil
 	case []byte:
 		b, _ := v.([]byte)
-		return tomlValueStringRepresentation(string(b), indent, arraysOneElementPerLine)
+		return tomlValueStringRepresentation(string(b), commented, indent, arraysOneElementPerLine)
 	case bool:
 		if value {
 			return "true", nil
@@ -152,7 +153,7 @@ func tomlValueStringRepresentation(v interface{}, indent string, arraysOneElemen
 		var values []string
 		for i := 0; i < rv.Len(); i++ {
 			item := rv.Index(i).Interface()
-			itemRepr, err := tomlValueStringRepresentation(item, indent, arraysOneElementPerLine)
+			itemRepr, err := tomlValueStringRepresentation(item, commented, indent, arraysOneElementPerLine)
 			if err != nil {
 				return "", err
 			}
@@ -166,12 +167,12 @@ func tomlValueStringRepresentation(v interface{}, indent string, arraysOneElemen
 
 			for _, value := range values {
 				stringBuffer.WriteString(valueIndent)
-				stringBuffer.WriteString(value)
+				stringBuffer.WriteString(commented + value)
 				stringBuffer.WriteString(`,`)
 				stringBuffer.WriteString("\n")
 			}
 
-			stringBuffer.WriteString(indent + "]")
+			stringBuffer.WriteString(indent + commented + "]")
 
 			return stringBuffer.String(), nil
 		}
@@ -270,10 +271,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)
+	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) (int64, error) {
+func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount int64, arraysOneElementPerLine bool, ord marshalOrder, parentCommented bool) (int64, error) {
 	var orderedVals []sortNode
 
 	switch ord {
@@ -293,10 +294,6 @@ func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount i
 			if keyspace != "" {
 				combinedKey = keyspace + "." + combinedKey
 			}
-			var commented string
-			if t.commented {
-				commented = "# "
-			}
 
 			switch node := v.(type) {
 			// node has to be of those two types given how keys are sorted above
@@ -317,24 +314,33 @@ func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount i
 						return bytesCount, errc
 					}
 				}
+
+				var commented string
+				if parentCommented || t.commented || tv.commented {
+					commented = "# "
+				}
 				writtenBytesCount, err := writeStrings(w, "\n", indent, commented, "[", combinedKey, "]\n")
 				bytesCount += int64(writtenBytesCount)
 				if err != nil {
 					return bytesCount, err
 				}
-				bytesCount, err = node.writeToOrdered(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine, ord)
+				bytesCount, err = node.writeToOrdered(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine, ord, parentCommented || t.commented || tv.commented)
 				if err != nil {
 					return bytesCount, err
 				}
 			case []*Tree:
 				for _, subTree := range node {
+					var commented string
+					if parentCommented || t.commented || subTree.commented {
+						commented = "# "
+					}
 					writtenBytesCount, err := writeStrings(w, "\n", indent, commented, "[[", combinedKey, "]]\n")
 					bytesCount += int64(writtenBytesCount)
 					if err != nil {
 						return bytesCount, err
 					}
 
-					bytesCount, err = subTree.writeToOrdered(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine, ord)
+					bytesCount, err = subTree.writeToOrdered(w, indent+"  ", combinedKey, bytesCount, arraysOneElementPerLine, ord, parentCommented || t.commented || subTree.commented)
 					if err != nil {
 						return bytesCount, err
 					}
@@ -347,7 +353,11 @@ func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount i
 				return bytesCount, fmt.Errorf("invalid value type at %s: %T", k, t.values[k])
 			}
 
-			repr, err := tomlValueStringRepresentation(v, indent, arraysOneElementPerLine)
+			var commented string
+			if parentCommented || t.commented || v.commented {
+				commented = "# "
+			}
+			repr, err := tomlValueStringRepresentation(v, commented, indent, arraysOneElementPerLine)
 			if err != nil {
 				return bytesCount, err
 			}
@@ -365,10 +375,6 @@ func (t *Tree) writeToOrdered(w io.Writer, indent, keyspace string, bytesCount i
 				}
 			}
 
-			var commented string
-			if v.commented {
-				commented = "# "
-			}
 			quotedKey := quoteKeyIfNeeded(k)
 			writtenBytesCount, err := writeStrings(w, indent, commented, quotedKey, " = ", repr, "\n")
 			bytesCount += int64(writtenBytesCount)