Support CommentToMap option

decode.go

@@ -28,6 +28,7 @@ type Decoder struct {
 	referenceReaders     []io.Reader
 	anchorNodeMap        map[string]ast.Node
 	anchorValueMap       map[string]reflect.Value
+	toCommentMap         CommentMap
 	opts                 []DecodeOption
 	referenceFiles       []string
 	referenceDirs        []string
@@ -115,6 +116,7 @@ func (d *Decoder) mapKeyNodeToString(node ast.Node) string {
 }
 
 func (d *Decoder) setToMapValue(node ast.Node, m map[string]interface{}) {
+	d.setPathToCommentMap(node)
 	switch n := node.(type) {
 	case *ast.MappingValueNode:
 		if n.Key.Type() == ast.MergeKeyType {
@@ -149,7 +151,30 @@ func (d *Decoder) setToOrderedMapValue(node ast.Node, m *MapSlice) {
 	}
 }
 
+func (d *Decoder) setPathToCommentMap(node ast.Node) {
+	if d.toCommentMap == nil {
+		return
+	}
+	commentGroup := node.GetComment()
+	if commentGroup == nil {
+		return
+	}
+	texts := []string{}
+	for _, comment := range commentGroup.Comments {
+		texts = append(texts, comment.Token.Value)
+	}
+	if len(texts) == 0 {
+		return
+	}
+	if len(texts) == 1 {
+		d.toCommentMap[node.GetPath()] = LineComment(texts[0])
+	} else {
+		d.toCommentMap[node.GetPath()] = HeadComment(texts...)
+	}
+}
+
 func (d *Decoder) nodeToValue(node ast.Node) interface{} {
+	d.setPathToCommentMap(node)
 	switch n := node.(type) {
 	case *ast.NullNode:
 		return nil
@@ -1434,7 +1459,11 @@ func (d *Decoder) resolveReference() error {
 }
 
 func (d *Decoder) parse(bytes []byte) (*ast.File, error) {
-	f, err := parser.ParseBytes(bytes, 0)
+	var parseMode parser.Mode
+	if d.toCommentMap != nil {
+		parseMode = parser.ParseComments
+	}
+	f, err := parser.ParseBytes(bytes, parseMode)
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to parse yaml")
 	}

error.go

@@ -11,6 +11,7 @@ var (
 	ErrInvalidPathString          = xerrors.New("invalid path string")
 	ErrNotFoundNode               = xerrors.New("node not found")
 	ErrUnknownCommentPositionType = xerrors.New("unknown comment position type")
+	ErrInvalidCommentMapValue     = xerrors.New("invalid comment map value. it must be not nil value")
 )
 
 func ErrUnsupportedHeadPositionType(node ast.Node) error {

option.go

@@ -216,3 +216,14 @@ func WithComment(cm CommentMap) EncodeOption {
 		return nil
 	}
 }
+
+// CommentToMap apply the position and content of comments in a YAML document to a CommentMap.
+func CommentToMap(cm CommentMap) DecodeOption {
+	return func(d *Decoder) error {
+		if cm == nil {
+			return ErrInvalidCommentMapValue
+		}
+		d.toCommentMap = cm
+		return nil
+	}
+}

yaml_test.go

@@ -2,6 +2,7 @@ package yaml_test
 
 import (
 	"bytes"
+	"reflect"
 	"testing"
 
 	"github.com/goccy/go-yaml"
@@ -466,3 +467,118 @@ baz:
 		}
 	})
 }
+
+func Test_CommentToMapOption(t *testing.T) {
+	t.Run("line comment", func(t *testing.T) {
+		yml := `
+foo: aaa #foo comment
+bar: #bar comment
+  bbb: ccc #bbb comment
+baz:
+  x: 10 #x comment
+`
+		var (
+			v  interface{}
+			cm = yaml.CommentMap{}
+		)
+		if err := yaml.UnmarshalWithOptions([]byte(yml), &v, yaml.CommentToMap(cm)); err != nil {
+			t.Fatal(err)
+		}
+		expected := []struct {
+			path    string
+			comment *yaml.Comment
+		}{
+			{
+				path:    "$.foo",
+				comment: yaml.LineComment("foo comment"),
+			},
+			{
+				path:    "$.bar",
+				comment: yaml.LineComment("bar comment"),
+			},
+			{
+				path:    "$.bar.bbb",
+				comment: yaml.LineComment("bbb comment"),
+			},
+			{
+				path:    "$.baz.x",
+				comment: yaml.LineComment("x comment"),
+			},
+		}
+		for _, exp := range expected {
+			comment := cm[exp.path]
+			if comment == nil {
+				t.Fatalf("failed to get path %s", exp.path)
+			}
+			if !reflect.DeepEqual(exp.comment, comment) {
+				t.Fatalf("failed to get comment. expected:[%+v] but got:[%+v]", exp.comment, comment)
+			}
+		}
+	})
+	t.Run("head comment", func(t *testing.T) {
+		yml := `
+#foo comment
+#foo comment2
+foo: aaa
+#bar comment
+#bar comment2
+bar:
+  #bbb comment
+  #bbb comment2
+  bbb: ccc
+baz:
+  #x comment
+  #x comment2
+  x: 10
+`
+		var (
+			v  interface{}
+			cm = yaml.CommentMap{}
+		)
+		if err := yaml.UnmarshalWithOptions([]byte(yml), &v, yaml.CommentToMap(cm)); err != nil {
+			t.Fatal(err)
+		}
+		expected := []struct {
+			path    string
+			comment *yaml.Comment
+		}{
+			{
+				path: "$.foo",
+				comment: yaml.HeadComment(
+					"foo comment",
+					"foo comment2",
+				),
+			},
+			{
+				path: "$.bar",
+				comment: yaml.HeadComment(
+					"bar comment",
+					"bar comment2",
+				),
+			},
+			{
+				path: "$.bar.bbb",
+				comment: yaml.HeadComment(
+					"bbb comment",
+					"bbb comment2",
+				),
+			},
+			{
+				path: "$.baz.x",
+				comment: yaml.HeadComment(
+					"x comment",
+					"x comment2",
+				),
+			},
+		}
+		for _, exp := range expected {
+			comment := cm[exp.path]
+			if comment == nil {
+				t.Fatalf("failed to get path %s", exp.path)
+			}
+			if !reflect.DeepEqual(exp.comment, comment) {
+				t.Fatalf("failed to get comment. expected:[%+v] but got:[%+v]", exp.comment, comment)
+			}
+		}
+	})
+}