Repository: goccy/go-yaml
Pull request #238 (2021-07-20 05:31:20)

Support WithComment option
Add option for encoding with comment .

Please see the following test case .
https://github.com/goccy/go-yaml/blob/3c954f0f01dfb5fa49a6f15e838b49046ea62989/yaml_test.go#L386-L468

`yaml.WithComment` is a option for encoding with comment .
The position where you want to add a comment is represented by `YAMLPath`, and it is the key of `yaml.CommentMap` .
Also, you can select `Head` comment or `Line` comment as the comment type.


Patch:
diff --git a/ast/ast.go b/ast/ast.go
index 30315a0..6a1aea7 100644
--- a/ast/ast.go
+++ b/ast/ast.go
@@ -1203,10 +1203,7 @@ func (n *MappingNode) blockStyleString(commentMode bool) string {
 			}
 			spaceNum++
 		}
-		comment := fmt.Sprintf(
-			"%s",
-			n.Comment.StringWithSpace(spaceNum),
-		)
+		comment := n.Comment.StringWithSpace(spaceNum)
 		return fmt.Sprintf("%s\n%s", comment, mapText)
 	}
 	return mapText
@@ -1391,6 +1388,14 @@ func (n *MappingValueNode) toString() string {
 			n.Value.String(),
 		)
 	}
+	if m, ok := n.Value.(*MappingNode); ok && m.Comment != nil {
+		return fmt.Sprintf(
+			"%s%s: %s",
+			space,
+			n.Key.String(),
+			strings.TrimLeft(n.Value.String(), " "),
+		)
+	}
 	return fmt.Sprintf("%s%s:\n%s", space, n.Key.String(), n.Value.String())
 }
 
@@ -1959,6 +1964,77 @@ func (v *filterWalker) Visit(n Node) Visitor {
 	return v
 }
 
+type parentFinder struct {
+	target Node
+}
+
+func (f *parentFinder) walk(parent, node Node) Node {
+	if f.target == node {
+		return parent
+	}
+	switch n := node.(type) {
+	case *CommentNode:
+		return nil
+	case *NullNode:
+		return nil
+	case *IntegerNode:
+		return nil
+	case *FloatNode:
+		return nil
+	case *StringNode:
+		return nil
+	case *MergeKeyNode:
+		return nil
+	case *BoolNode:
+		return nil
+	case *InfinityNode:
+		return nil
+	case *NanNode:
+		return nil
+	case *LiteralNode:
+		return f.walk(node, n.Value)
+	case *DirectiveNode:
+		return f.walk(node, n.Value)
+	case *TagNode:
+		return f.walk(node, n.Value)
+	case *DocumentNode:
+		return f.walk(node, n.Body)
+	case *MappingNode:
+		for _, value := range n.Values {
+			if found := f.walk(node, value); found != nil {
+				return found
+			}
+		}
+	case *MappingKeyNode:
+		return f.walk(node, n.Value)
+	case *MappingValueNode:
+		if found := f.walk(node, n.Key); found != nil {
+			return found
+		}
+		return f.walk(node, n.Value)
+	case *SequenceNode:
+		for _, value := range n.Values {
+			if found := f.walk(node, value); found != nil {
+				return found
+			}
+		}
+	case *AnchorNode:
+		if found := f.walk(node, n.Name); found != nil {
+			return found
+		}
+		return f.walk(node, n.Value)
+	case *AliasNode:
+		return f.walk(node, n.Value)
+	}
+	return nil
+}
+
+// Parent get parent node from child node.
+func Parent(root, child Node) Node {
+	finder := &parentFinder{target: child}
+	return finder.walk(root, root)
+}
+
 // Filter returns a list of nodes that match the given type.
 func Filter(typ NodeType, node Node) []Node {
 	walker := &filterWalker{typ: typ}
diff --git a/encode.go b/encode.go
index aa63a5d..8ed0694 100644
--- a/encode.go
+++ b/encode.go
@@ -37,6 +37,7 @@ type Encoder struct {
 	anchorCallback             func(*ast.AnchorNode, interface{}) error
 	anchorPtrToNameMap         map[uintptr]string
 	useLiteralStyleIfMultiline bool
+	commentMap                 map[*Path]*Comment
 
 	line        int
 	column      int
@@ -81,6 +82,9 @@ func (e *Encoder) EncodeContext(ctx context.Context, v interface{}) error {
 	if err != nil {
 		return errors.Wrapf(err, "failed to encode to node")
 	}
+	if err := e.setCommentByCommentMap(node); err != nil {
+		return errors.Wrapf(err, "failed to set comment by comment map")
+	}
 	var p printer.Printer
 	e.writer.Write(p.PrintNode(node))
 	return nil
@@ -105,6 +109,49 @@ func (e *Encoder) EncodeToNodeContext(ctx context.Context, v interface{}) (ast.N
 	return node, nil
 }
 
+func (e *Encoder) setCommentByCommentMap(node ast.Node) error {
+	if e.commentMap == nil {
+		return nil
+	}
+	for path, comment := range e.commentMap {
+		n, err := path.FilterNode(node)
+		if err != nil {
+			return errors.Wrapf(err, "failed to filter node")
+		}
+		comments := []*token.Token{}
+		for _, text := range comment.Texts {
+			comments = append(comments, token.New(text, text, nil))
+		}
+		commentGroup := ast.CommentGroup(comments)
+		switch comment.Position {
+		case CommentLinePosition:
+			if err := n.SetComment(commentGroup); err != nil {
+				return errors.Wrapf(err, "failed to set comment")
+			}
+		case CommentHeadPosition:
+			parent := ast.Parent(node, n)
+			if parent == nil {
+				return ErrUnsupportedHeadPositionType(node)
+			}
+			switch node := parent.(type) {
+			case *ast.MappingValueNode:
+				if err := node.SetComment(commentGroup); err != nil {
+					return errors.Wrapf(err, "failed to set comment")
+				}
+			case *ast.MappingNode:
+				if err := node.SetComment(commentGroup); err != nil {
+					return errors.Wrapf(err, "failed to set comment")
+				}
+			default:
+				return ErrUnsupportedHeadPositionType(node)
+			}
+		default:
+			return ErrUnknownCommentPositionType
+		}
+	}
+	return nil
+}
+
 func (e *Encoder) encodeDocument(doc []byte) (ast.Node, error) {
 	f, err := parser.ParseBytes(doc, 0)
 	if err != nil {
diff --git a/error.go b/error.go
index 3dfaf21..61c6eff 100644
--- a/error.go
+++ b/error.go
@@ -6,12 +6,17 @@ import (
 )
 
 var (
-	ErrInvalidQuery      = xerrors.New("invalid query")
-	ErrInvalidPath       = xerrors.New("invalid path instance")
-	ErrInvalidPathString = xerrors.New("invalid path string")
-	ErrNotFoundNode      = xerrors.New("node not found")
+	ErrInvalidQuery               = xerrors.New("invalid query")
+	ErrInvalidPath                = xerrors.New("invalid path instance")
+	ErrInvalidPathString          = xerrors.New("invalid path string")
+	ErrNotFoundNode               = xerrors.New("node not found")
+	ErrUnknownCommentPositionType = xerrors.New("unknown comment position type")
 )
 
+func ErrUnsupportedHeadPositionType(node ast.Node) error {
+	return xerrors.Errorf("unsupported comment head position for %s", node.Type())
+}
+
 // IsInvalidQueryError whether err is ErrInvalidQuery or not.
 func IsInvalidQueryError(err error) bool {
 	return xerrors.Is(err, ErrInvalidQuery)
diff --git a/option.go b/option.go
index 5dc72c5..8c8f75d 100644
--- a/option.go
+++ b/option.go
@@ -156,3 +156,63 @@ func UseJSONMarshaler() EncodeOption {
 		return nil
 	}
 }
+
+// CommentPosition type of the position for comment.
+type CommentPosition int
+
+const (
+	CommentLinePosition CommentPosition = iota
+	CommentHeadPosition
+)
+
+func (p CommentPosition) String() string {
+	switch p {
+	case CommentLinePosition:
+		return "Line"
+	case CommentHeadPosition:
+		return "Head"
+	default:
+		return ""
+	}
+}
+
+// LineComment create a one-line comment for CommentMap.
+func LineComment(text string) *Comment {
+	return &Comment{
+		Texts:    []string{text},
+		Position: CommentLinePosition,
+	}
+}
+
+// HeadComment create a multiline comment for CommentMap.
+func HeadComment(texts ...string) *Comment {
+	return &Comment{
+		Texts:    texts,
+		Position: CommentHeadPosition,
+	}
+}
+
+// Comment raw data for comment.
+type Comment struct {
+	Texts    []string
+	Position CommentPosition
+}
+
+// CommentMap map of the position of the comment and the comment information.
+type CommentMap map[string]*Comment
+
+// WithComment add a comment using the location and text information given in the CommentMap.
+func WithComment(cm CommentMap) EncodeOption {
+	return func(e *Encoder) error {
+		commentMap := map[*Path]*Comment{}
+		for k, v := range cm {
+			path, err := PathString(k)
+			if err != nil {
+				return err
+			}
+			commentMap[path] = v
+		}
+		e.commentMap = commentMap
+		return nil
+	}
+}


Test patch:
diff --git a/yaml_test.go b/yaml_test.go
index 8141dd2..df588f7 100644
--- a/yaml_test.go
+++ b/yaml_test.go
@@ -382,3 +382,87 @@ a: 1
 		t.Fatalf("failed to convert json to yaml: expected [%q] but got [%q]", expected, actual)
 	}
 }
+
+func Test_WithCommentOption(t *testing.T) {
+	v := struct {
+		Foo string                 `yaml:"foo"`
+		Bar map[string]interface{} `yaml:"bar"`
+		Baz struct {
+			X int `yaml:"x"`
+		} `yaml:"baz"`
+	}{
+		Foo: "aaa",
+		Bar: map[string]interface{}{"bbb": "ccc"},
+		Baz: struct {
+			X int `yaml:"x"`
+		}{X: 10},
+	}
+	t.Run("line comment", func(t *testing.T) {
+		b, err := yaml.MarshalWithOptions(v, yaml.WithComment(
+			yaml.CommentMap{
+				"$.foo":     yaml.LineComment("foo comment"),
+				"$.bar":     yaml.LineComment("bar comment"),
+				"$.bar.bbb": yaml.LineComment("bbb comment"),
+				"$.baz.x":   yaml.LineComment("x comment"),
+			},
+		))
+		if err != nil {
+			t.Fatal(err)
+		}
+		expected := `
+foo: aaa #foo comment
+bar: #bar comment
+  bbb: ccc #bbb comment
+baz:
+  x: 10 #x comment
+`
+		actual := "\n" + string(b)
+		if expected != actual {
+			t.Fatalf("expected:%s but got %s", expected, actual)
+		}
+	})
+	t.Run("head comment", func(t *testing.T) {
+		b, err := yaml.MarshalWithOptions(v, yaml.WithComment(
+			yaml.CommentMap{
+				"$.foo": yaml.HeadComment(
+					"foo comment",
+					"foo comment2",
+				),
+				"$.bar": yaml.HeadComment(
+					"bar comment",
+					"bar comment2",
+				),
+				"$.bar.bbb": yaml.HeadComment(
+					"bbb comment",
+					"bbb comment2",
+				),
+				"$.baz.x": yaml.HeadComment(
+					"x comment",
+					"x comment2",
+				),
+			},
+		))
+		if err != nil {
+			t.Fatal(err)
+		}
+		expected := `
+#foo comment
+#foo comment2
+foo: aaa
+#bar comment
+#bar comment2
+bar:
+  #bbb comment
+  #bbb comment2
+  bbb: ccc
+baz:
+  #x comment
+  #x comment2
+  x: 10
+`
+		actual := "\n" + string(b)
+		if expected != actual {
+			t.Fatalf("expected:%s but got %s", expected, actual)
+		}
+	})
+}