Comment annotation for Marshal (#185)

marshal.go

@@ -5,12 +5,15 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"strconv"
 	"strings"
 	"time"
 )
 
 type tomlOpts struct {
 	name      string
+	comment   string
+	commented bool
 	include   bool
 	omitempty bool
 }
@@ -147,7 +150,7 @@ func valueToTree(mtype reflect.Type, mval reflect.Value) (*Tree, error) {
 				if err != nil {
 					return nil, err
 				}
-				tval.Set(opts.name, val)
+				tval.Set(opts.name, opts.comment, opts.commented, val)
 			}
 		}
 	case reflect.Map:
@@ -157,7 +160,7 @@ func valueToTree(mtype reflect.Type, mval reflect.Value) (*Tree, error) {
 			if err != nil {
 				return nil, err
 			}
-			tval.Set(key.String(), val)
+			tval.Set(key.String(), "", false, val)
 		}
 	}
 	return tval, nil
@@ -448,7 +451,12 @@ func unwrapPointer(mtype reflect.Type, tval interface{}) (reflect.Value, error)
 func tomlOptions(vf reflect.StructField) tomlOpts {
 	tag := vf.Tag.Get("toml")
 	parse := strings.Split(tag, ",")
-	result := tomlOpts{vf.Name, true, false}
+	var comment string
+	if c := vf.Tag.Get("comment"); c != "" {
+		comment = c
+	}
+	commented, _ := strconv.ParseBool(vf.Tag.Get("commented"))
+	result := tomlOpts{name: vf.Name, comment: comment, commented: commented, include: true, omitempty: false}
 	if parse[0] != "" {
 		if parse[0] == "-" && len(parse) == 1 {
 			result.include = false

marshal_test.go

@@ -598,3 +598,54 @@ func TestNestedCustomMarshaler(t *testing.T) {
 		t.Errorf("Bad nested custom marshaler: expected\n-----\n%s\n-----\ngot\n-----\n%s\n-----\n", expected, result)
 	}
 }
+
+var commentTestToml = []byte(`
+# it's a comment on type
+[postgres]
+  # isCommented = "dvalue"
+  noComment = "cvalue"
+
+  # A comment on AttrB with a
+  # break line
+  password = "bvalue"
+
+  # A comment on AttrA
+  user = "avalue"
+
+  [[postgres.My]]
+
+    # a comment on my on typeC
+    My = "Foo"
+
+  [[postgres.My]]
+
+    # a comment on my on typeC
+    My = "Baar"
+`)
+
+func TestMarshalComment(t *testing.T) {
+	type TypeC struct {
+		My string `comment:"a comment on my on typeC"`
+	}
+	type TypeB struct {
+		AttrA string `toml:"user" comment:"A comment on AttrA"`
+		AttrB string `toml:"password" comment:"A comment on AttrB with a\n break line"`
+		AttrC string `toml:"noComment"`
+		AttrD string `toml:"isCommented" commented:"true"`
+		My    []TypeC
+	}
+	type TypeA struct {
+		TypeB TypeB `toml:"postgres" comment:"it's a comment on type"`
+	}
+
+	ta := []TypeC{{My: "Foo"}, {My: "Baar"}}
+	config := TypeA{TypeB{AttrA: "avalue", AttrB: "bvalue", AttrC: "cvalue", AttrD: "dvalue", My: ta}}
+	result, err := Marshal(config)
+	if err != nil {
+		t.Fatal(err)
+	}
+	expected := commentTestToml
+	if !bytes.Equal(result, expected) {
+		t.Errorf("Bad marshal: expected\n-----\n%s\n-----\ngot\n-----\n%s\n-----\n", expected, result)
+	}
+}

parser.go

@@ -110,7 +110,7 @@ func (p *tomlParser) parseGroupArray() tomlParserStateFn {
 	newTree := newTree()
 	newTree.position = startToken.Position
 	array = append(array, newTree)
-	p.tree.SetPath(p.currentTable, array)
+	p.tree.SetPath(p.currentTable, "", false, array)
 
 	// remove all keys that were children of this table array
 	prefix := key.val + "."
@@ -205,7 +205,7 @@ func (p *tomlParser) parseAssign() tomlParserStateFn {
 	case *Tree, []*Tree:
 		toInsert = value
 	default:
-		toInsert = &tomlValue{value, key.Position}
+		toInsert = &tomlValue{value: value, position: key.Position}
 	}
 	targetNode.values[keyVal] = toInsert
 	return p.parseStart
@@ -299,7 +299,7 @@ Loop:
 			key := p.getToken()
 			p.assume(tokenEqual)
 			value := p.parseRvalue()
-			tree.Set(key.val, value)
+			tree.Set(key.val, "", false, value)
 		case tokenComma:
 			if previous == nil {
 				p.raiseError(follow, "inline table cannot start with a comma")

parser_test.go

@@ -46,7 +46,7 @@ func assertTree(t *testing.T, tree *Tree, err error, ref map[string]interface{})
 func TestCreateSubTree(t *testing.T) {
 	tree := newTree()
 	tree.createSubTree([]string{"a", "b", "c"}, Position{})
-	tree.Set("a.b.c", 42)
+	tree.Set("a.b.c", "", false, 42)
 	if tree.Get("a.b.c") != 42 {
 		t.Fail()
 	}

toml.go

@@ -11,14 +11,18 @@ import (
 )
 
 type tomlValue struct {
-	value    interface{} // string, int64, uint64, float64, bool, time.Time, [] of any of this list
+	value     interface{} // string, int64, uint64, float64, bool, time.Time, [] of any of this list
-	position Position
+	comment   string
+	commented bool
+	position  Position
 }
 
 // Tree is the result of the parsing of a TOML file.
 type Tree struct {
-	values   map[string]interface{} // string -> *tomlValue, *Tree, []*Tree
+	values    map[string]interface{} // string -> *tomlValue, *Tree, []*Tree
-	position Position
+	comment   string
+	commented bool
+	position  Position
 }
 
 func newTree() *Tree {
@@ -177,14 +181,14 @@ func (t *Tree) GetDefault(key string, def interface{}) interface{} {
 // Set an element in the tree.
 // Key is a dot-separated path (e.g. a.b.c).
 // Creates all necessary intermediate trees, if needed.
-func (t *Tree) Set(key string, value interface{}) {
+func (t *Tree) Set(key string, comment string, commented bool, value interface{}) {
-	t.SetPath(strings.Split(key, "."), value)
+	t.SetPath(strings.Split(key, "."), comment, commented, value)
 }
 
 // SetPath sets an element in the tree.
 // Keys is an array of path elements (e.g. {"a","b","c"}).
 // Creates all necessary intermediate trees, if needed.
-func (t *Tree) SetPath(keys []string, value interface{}) {
+func (t *Tree) SetPath(keys []string, comment string, commented bool, value interface{}) {
 	subtree := t
 	for _, intermediateKey := range keys[:len(keys)-1] {
 		nextTree, exists := subtree.values[intermediateKey]
@@ -209,13 +213,17 @@ func (t *Tree) SetPath(keys []string, value interface{}) {
 
 	switch value.(type) {
 	case *Tree:
+		tt := value.(*Tree)
+		tt.comment = comment
 		toInsert = value
 	case []*Tree:
 		toInsert = value
 	case *tomlValue:
-		toInsert = value
+		tt := value.(*tomlValue)
+		tt.comment = comment
+		toInsert = tt
 	default:
-		toInsert = &tomlValue{value: value}
+		toInsert = &tomlValue{value: value, comment: comment, commented: commented}
 	}
 
 	subtree.values[keys[len(keys)-1]] = toInsert

tomltree_create.go

@@ -104,7 +104,7 @@ func sliceToTree(object interface{}) (interface{}, error) {
 		}
 		arrayValue = reflect.Append(arrayValue, reflect.ValueOf(simpleValue))
 	}
-	return &tomlValue{arrayValue.Interface(), Position{}}, nil
+	return &tomlValue{value: arrayValue.Interface(), position: Position{}}, nil
 }
 
 func toTree(object interface{}) (interface{}, error) {
@@ -127,7 +127,7 @@ func toTree(object interface{}) (interface{}, error) {
 			}
 			values[key.String()] = newValue
 		}
-		return &Tree{values, Position{}}, nil
+		return &Tree{values: values, position: Position{}}, nil
 	}
 
 	if value.Kind() == reflect.Array || value.Kind() == reflect.Slice {
@@ -138,5 +138,5 @@ func toTree(object interface{}) (interface{}, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &tomlValue{simpleValue, Position{}}, nil
+	return &tomlValue{value: simpleValue, position: Position{}}, nil
 }

tomltree_write.go

@@ -118,7 +118,24 @@ func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64) (
 			return bytesCount, err
 		}
 
-		writtenBytesCount, err := writeStrings(w, indent, k, " = ", repr, "\n")
+		if v.comment != "" {
+			comment := strings.Replace(v.comment, "\n", "\n"+indent+"#", -1)
+			start := "# "
+			if strings.HasPrefix(comment, "#") {
+				start = ""
+			}
+			writtenBytesCountComment, errc := writeStrings(w, "\n", indent, start, comment, "\n")
+			bytesCount += int64(writtenBytesCountComment)
+			if errc != nil {
+				return bytesCount, errc
+			}
+		}
+
+		var commented string
+		if v.commented {
+			commented = "# "
+		}
+		writtenBytesCount, err := writeStrings(w, indent, commented, k, " = ", repr, "\n")
 		bytesCount += int64(writtenBytesCount)
 		if err != nil {
 			return bytesCount, err
@@ -132,11 +149,31 @@ func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64) (
 		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
 		case *Tree:
-			writtenBytesCount, err := writeStrings(w, "\n", indent, "[", combinedKey, "]\n")
+			tv, ok := t.values[k].(*Tree)
+			if !ok {
+				return bytesCount, fmt.Errorf("invalid value type at %s: %T", k, t.values[k])
+			}
+			if tv.comment != "" {
+				comment := strings.Replace(tv.comment, "\n", "\n"+indent+"#", -1)
+				start := "# "
+				if strings.HasPrefix(comment, "#") {
+					start = ""
+				}
+				writtenBytesCountComment, errc := writeStrings(w, "\n", indent, start, comment)
+				bytesCount += int64(writtenBytesCountComment)
+				if errc != nil {
+					return bytesCount, errc
+				}
+			}
+			writtenBytesCount, err := writeStrings(w, "\n", indent, commented, "[", combinedKey, "]\n")
 			bytesCount += int64(writtenBytesCount)
 			if err != nil {
 				return bytesCount, err
@@ -147,7 +184,7 @@ func (t *Tree) writeTo(w io.Writer, indent, keyspace string, bytesCount int64) (
 			}
 		case []*Tree:
 			for _, subTree := range node {
-				writtenBytesCount, err := writeStrings(w, "\n", indent, "[[", combinedKey, "]]\n")
+				writtenBytesCount, err := writeStrings(w, "\n", indent, commented, "[[", combinedKey, "]]\n")
 				bytesCount += int64(writtenBytesCount)
 				if err != nil {
 					return bytesCount, err

tomltree_write_test.go

@@ -161,13 +161,13 @@ func TestTreeWriteToInvalidTreeSimpleValue(t *testing.T) {
 }
 
 func TestTreeWriteToInvalidTreeTomlValue(t *testing.T) {
-	tree := Tree{values: map[string]interface{}{"foo": &tomlValue{int8(1), Position{}}}}
+	tree := Tree{values: map[string]interface{}{"foo": &tomlValue{value: int8(1), comment: "", position: Position{}}}}
 	_, err := tree.ToTomlString()
 	assertErrorString(t, "unsupported value type int8: 1", err)
 }
 
 func TestTreeWriteToInvalidTreeTomlValueArray(t *testing.T) {
-	tree := Tree{values: map[string]interface{}{"foo": &tomlValue{[]interface{}{int8(1)}, Position{}}}}
+	tree := Tree{values: map[string]interface{}{"foo": &tomlValue{value: int8(1), comment: "", position: Position{}}}}
 	_, err := tree.ToTomlString()
 	assertErrorString(t, "unsupported value type int8: 1", err)
 }