Fix comment parsing with null value

parser/parser.go

@@ -156,15 +156,34 @@ func (p *parser) createMapValueNode(ctx *context, key ast.MapKeyNode, colonToken
 		ctx.insertToken(ctx.idx, nullToken)
 		return ast.Null(nullToken), nil
 	}
-
+	var comment *ast.CommentGroupNode
+	if tk.Type == token.CommentType {
+		comment = p.parseCommentOnly(ctx)
+		if comment != nil {
+			comment.SetPath(ctx.withChild(key.GetToken().Value).path)
+		}
+		tk = ctx.currentToken()
+	}
 	if tk.Position.Column == key.GetToken().Position.Column && tk.Type == token.StringType {
 		// in this case,
 		// ----
 		// key: <value does not defined>
 		// next
+
 		nullToken := p.createNullToken(colonToken)
 		ctx.insertToken(ctx.idx, nullToken)
-		return ast.Null(nullToken), nil
+		nullNode := ast.Null(nullToken)
+
+		// If there is a comment, it is already bound to the key node,
+		// so remove the comment from the key to bind it to the null value.
+		comment := key.GetComment()
+		if comment != nil {
+			if err := key.SetComment(nil); err != nil {
+				return nil, err
+			}
+			nullNode.SetComment(comment)
+		}
+		return nullNode, nil
 	}
 
 	if tk.Position.Column < key.GetToken().Position.Column {
@@ -174,13 +193,20 @@ func (p *parser) createMapValueNode(ctx *context, key ast.MapKeyNode, colonToken
 		// next
 		nullToken := p.createNullToken(colonToken)
 		ctx.insertToken(ctx.idx, nullToken)
-		return ast.Null(nullToken), nil
+		nullNode := ast.Null(nullToken)
+		if comment != nil {
+			nullNode.SetComment(comment)
+		}
+		return nullNode, nil
 	}
 
 	value, err := p.parseToken(ctx, ctx.currentToken())
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to parse mapping 'value' node")
 	}
+	if comment != nil {
+		value.SetComment(comment)
+	}
 	return value, nil
 }
 

parser/parser_test.go

@@ -808,6 +808,49 @@ foo: > # comment
 	}
 }
 
+func TestCommentWithNull(t *testing.T) {
+	expected := `
+foo:
+  bar: null # comment
+  baz: 1
+`
+	t.Run("same line", func(t *testing.T) {
+		content := `
+foo:
+  bar: # comment
+  baz: 1
+`
+		f, err := parser.ParseBytes([]byte(content), parser.ParseComments)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(f.Docs) != 1 {
+			t.Fatal("failed to parse content with same line comment")
+		}
+		if f.Docs[0].String()+"\n" != strings.TrimPrefix(expected, "\n") {
+			t.Fatal("failed to parse comment")
+		}
+	})
+	t.Run("next line", func(t *testing.T) {
+		content := `
+foo:
+  bar:
+    # comment
+  baz: 1
+`
+		f, err := parser.ParseBytes([]byte(content), parser.ParseComments)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(f.Docs) != 1 {
+			t.Fatal("failed to parse content with next line comment")
+		}
+		if f.Docs[0].String()+"\n" != strings.TrimPrefix(expected, "\n") {
+			t.Fatal("failed to parse comment")
+		}
+	})
+}
+
 func TestNodePath(t *testing.T) {
 	yml := `
 a: # commentA