Repository: goccy/go-yaml
Pull request #234 (2021-07-19 09:49:21)

Fix parsing of literal with comment
fix https://github.com/goccy/go-yaml/issues/197




Patch:
diff --git a/ast/ast.go b/ast/ast.go
index b035b36..30315a0 100644
--- a/ast/ast.go
+++ b/ast/ast.go
@@ -879,7 +879,11 @@ func (n *LiteralNode) GetValue() interface{} {
 // String literal to text
 func (n *LiteralNode) String() string {
 	origin := n.Value.GetToken().Origin
-	return fmt.Sprintf("%s\n%s", n.Start.Value, strings.TrimRight(strings.TrimRight(origin, " "), "\n"))
+	lit := strings.TrimRight(strings.TrimRight(origin, " "), "\n")
+	if n.Comment != nil {
+		return fmt.Sprintf("%s %s\n%s", n.Start.Value, n.Comment.String(), lit)
+	}
+	return fmt.Sprintf("%s\n%s", n.Start.Value, lit)
 }
 
 func (n *LiteralNode) stringWithoutComment() string {
diff --git a/parser/parser.go b/parser/parser.go
index 1b646d0..70f6e5c 100644
--- a/parser/parser.go
+++ b/parser/parser.go
@@ -423,7 +423,17 @@ func (p *parser) parseDirective(ctx *context) (ast.Node, error) {
 func (p *parser) parseLiteral(ctx *context) (ast.Node, error) {
 	node := ast.Literal(ctx.currentToken())
 	ctx.progress(1) // skip literal/folded token
-	value, err := p.parseToken(ctx, ctx.currentToken())
+
+	tk := ctx.currentToken()
+	var comment *ast.CommentGroupNode
+	if tk.Type == token.CommentType {
+		comment = p.parseCommentOnly(ctx)
+		if err := node.SetComment(comment); err != nil {
+			return nil, errors.Wrapf(err, "failed to set comment to literal")
+		}
+		tk = ctx.currentToken()
+	}
+	value, err := p.parseToken(ctx, tk)
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to parse literal/folded value")
 	}
diff --git a/scanner/scanner.go b/scanner/scanner.go
index bb85d89..7082c5a 100644
--- a/scanner/scanner.go
+++ b/scanner/scanner.go
@@ -456,6 +456,17 @@ func (s *Scanner) scanComment(ctx *Context) (tk *token.Token, pos int) {
 	return
 }
 
+func trimCommentFromLiteralOpt(text string) (string, error) {
+	idx := strings.Index(text, "#")
+	if idx < 0 {
+		return text, nil
+	}
+	if idx == 0 {
+		return "", xerrors.New("invalid literal header")
+	}
+	return text[:idx-1], nil
+}
+
 func (s *Scanner) scanLiteral(ctx *Context, c rune) {
 	ctx.addOriginBuf(c)
 	if ctx.isEOS() {
@@ -498,14 +509,44 @@ func (s *Scanner) scanLiteralHeader(ctx *Context) (pos int, err error) {
 		case '\n', '\r':
 			value := ctx.source(ctx.idx, ctx.idx+idx)
 			opt := strings.TrimRight(value, " ")
+			orgOptLen := len(opt)
+			opt, err = trimCommentFromLiteralOpt(opt)
+			if err != nil {
+				return
+			}
 			switch opt {
 			case "", "+", "-",
 				"0", "1", "2", "3", "4", "5", "6", "7", "8", "9":
+				hasComment := len(opt) < orgOptLen
 				if header == '|' {
-					ctx.addToken(token.Literal("|"+opt, string(ctx.obuf), s.pos()))
+					if hasComment {
+						commentLen := orgOptLen - len(opt)
+						headerPos := strings.Index(string(ctx.obuf), "|")
+						litBuf := ctx.obuf[:len(ctx.obuf)-commentLen-headerPos]
+						commentBuf := ctx.obuf[len(litBuf):]
+						ctx.addToken(token.Literal("|"+opt, string(litBuf), s.pos()))
+						s.column += len(litBuf)
+						s.offset += len(litBuf)
+						commentHeader := strings.Index(value, "#")
+						ctx.addToken(token.Comment(string(value[commentHeader+1:]), string(commentBuf), s.pos()))
+					} else {
+						ctx.addToken(token.Literal("|"+opt, string(ctx.obuf), s.pos()))
+					}
 					ctx.isLiteral = true
 				} else if header == '>' {
-					ctx.addToken(token.Folded(">"+opt, string(ctx.obuf), s.pos()))
+					if hasComment {
+						commentLen := orgOptLen - len(opt)
+						headerPos := strings.Index(string(ctx.obuf), ">")
+						foldedBuf := ctx.obuf[:len(ctx.obuf)-commentLen-headerPos]
+						commentBuf := ctx.obuf[len(foldedBuf):]
+						ctx.addToken(token.Folded(">"+opt, string(foldedBuf), s.pos()))
+						s.column += len(foldedBuf)
+						s.offset += len(foldedBuf)
+						commentHeader := strings.Index(value, "#")
+						ctx.addToken(token.Comment(string(value[commentHeader+1:]), string(commentBuf), s.pos()))
+					} else {
+						ctx.addToken(token.Folded(">"+opt, string(ctx.obuf), s.pos()))
+					}
 					ctx.isFolded = true
 				}
 				s.indentState = IndentStateKeep


Test patch:
diff --git a/parser/parser_test.go b/parser/parser_test.go
index c579ccd..f6a97e6 100644
--- a/parser/parser_test.go
+++ b/parser/parser_test.go
@@ -727,6 +727,20 @@ foo2: &anchor text # anchor comment
 # foo3 comment
 # foo3 comment2
 foo3: *anchor # alias comment
+`,
+		},
+		{
+			name: "literal",
+			yaml: `
+foo: | # comment
+  x: 42
+`,
+		},
+		{
+			name: "folded",
+			yaml: `
+foo: > # comment
+  x: 42
 `,
 		},
 	}