Merge pull request #328 from goccy/fix-quoted-map-key

Fix quoted map key

decode_test.go

@@ -2144,6 +2144,29 @@ b: *a
 			t.Fatal("failed to unmarshal")
 		}
 	})
+	t.Run("quoted map keys", func(t *testing.T) {
+		t.Parallel()
+		yml := `
+a:
+  "b"  : 2
+  'c': true
+`
+		var v struct {
+			A struct {
+				B int
+				C bool
+			}
+		}
+		if err := yaml.Unmarshal([]byte(yml), &v); err != nil {
+			t.Fatalf("failed to unmarshal %v", err)
+		}
+		if v.A.B != 2 {
+			t.Fatalf("expected a.b to equal 2 but was %d", v.A.B)
+		}
+		if !v.A.C {
+			t.Fatal("expected a.c to be true but was false")
+		}
+	})
 }
 
 type unmarshalablePtrStringContainer struct {

lexer/lexer_test.go

@@ -56,6 +56,10 @@ func TestTokenize(t *testing.T) {
 		"a: 'Hello #comment'\n",
 		"a: 100.5\n",
 		"a: bogus\n",
+		"\"a\": double quoted map key",
+		"'a': single quoted map key",
+		"a: \"double quoted\"\nb: \"value map\"",
+		"a: 'single quoted'\nb: 'value map'",
 	}
 	for _, src := range sources {
 		lexer.Tokenize(src).Dump()
@@ -231,6 +235,60 @@ func TestSingleLineToken_ValueLineColumnPosition(t *testing.T) {
 				15: "]",
 			},
 		},
+		{
+			name: "double quote key",
+			src:  `"a": b`,
+			expect: map[int]string{
+				1: "a",
+				4: ":",
+				6: "b",
+			},
+		},
+		{
+			name: "single quote key",
+			src:  `'a': b`,
+			expect: map[int]string{
+				1: "a",
+				4: ":",
+				6: "b",
+			},
+		},
+		{
+			name: "double quote key and value",
+			src:  `"a": "b"`,
+			expect: map[int]string{
+				1: "a",
+				4: ":",
+				6: "b",
+			},
+		},
+		{
+			name: "single quote key and value",
+			src:  `'a': 'b'`,
+			expect: map[int]string{
+				1: "a",
+				4: ":",
+				6: "b",
+			},
+		},
+		{
+			name: "double quote key, single quote value",
+			src:  `"a": 'b'`,
+			expect: map[int]string{
+				1: "a",
+				4: ":",
+				6: "b",
+			},
+		},
+		{
+			name: "single quote key, double quote value",
+			src:  `'a': "b"`,
+			expect: map[int]string{
+				1: "a",
+				4: ":",
+				6: "b",
+			},
+		},
 	}
 
 	for _, tc := range tests {
@@ -432,6 +490,59 @@ foo2: 'bar2'`,
 				},
 			},
 		},
+		{
+			name: "single and double quote map keys",
+			src: `"a": test
+'b': 1
+c: true`,
+			expect: []testToken{
+				{
+					line:   1,
+					column: 1,
+					value:  "a",
+				},
+				{
+					line:   1,
+					column: 4,
+					value:  ":",
+				},
+				{
+					line:   1,
+					column: 6,
+					value:  "test",
+				},
+				{
+					line:   2,
+					column: 1,
+					value:  "b",
+				},
+				{
+					line:   2,
+					column: 4,
+					value:  ":",
+				},
+				{
+					line:   2,
+					column: 6,
+					value:  "1",
+				},
+				{
+					line:   3,
+					column: 1,
+					value:  "c",
+				},
+				{
+					line:   3,
+					column: 2,
+					value:  ":",
+				},
+				{
+					line:   3,
+					column: 4,
+					value:  "true",
+				},
+			},
+		},
 	}
 
 	for _, tc := range tests {

parser/parser_test.go

@@ -80,6 +80,8 @@ func TestParser(t *testing.T) {
   ? !!str "implicit" : !!str "entry",
   ? !!null "" : !!null "",
 }`,
+		"\"a\": a\n\"b\": b",
+		"'a': a\n'b': b",
 	}
 	for _, src := range sources {
 		if _, err := parser.Parse(lexer.Tokenize(src), 0); err != nil {
@@ -562,6 +564,22 @@ b: c
 			`
 - key1: val
   key2: ( foo + bar )
+`,
+		},
+		{
+			`
+"a": b
+'c': d
+"e": "f"
+g: "h"
+i: 'j'
+`,
+			`
+"a": b
+'c': d
+"e": "f"
+g: "h"
+i: 'j'
 `,
 		},
 	}

scanner/context.go

@@ -6,6 +6,8 @@ import (
 	"github.com/goccy/go-yaml/token"
 )
 
+const whitespace = ' '
+
 // Context context at scanning
 type Context struct {
 	idx                int
@@ -143,7 +145,22 @@ func (c *Context) previousChar() rune {
 }
 
 func (c *Context) currentChar() rune {
-	return c.src[c.idx]
+	if c.size > c.idx {
+		return c.src[c.idx]
+	}
+	return rune(0)
+}
+
+func (c *Context) currentCharWithSkipWhitespace() rune {
+	idx := c.idx
+	for c.size > idx {
+		ch := c.src[idx]
+		if ch != whitespace {
+			return ch
+		}
+		idx++
+	}
+	return rune(0)
 }
 
 func (c *Context) nextChar() rune {
@@ -203,3 +220,10 @@ func (c *Context) bufferedToken(pos *token.Position) *token.Token {
 	c.resetBuffer()
 	return tk
 }
+
+func (c *Context) lastToken() *token.Token {
+	if len(c.tokens) != 0 {
+		return c.tokens[len(c.tokens)-1]
+	}
+	return nil
+}

scanner/scanner.go

@@ -733,6 +733,12 @@ func (s *Scanner) scan(ctx *Context) (pos int) {
 				if tk != nil {
 					s.prevIndentColumn = tk.Position.Column
 					ctx.addToken(tk)
+				} else if tk := ctx.lastToken(); tk != nil {
+					// If the map key is quote, the buffer does not exist because it has already been cut into tokens.
+					// Therefore, we need to check the last token.
+					if tk.Indicator == token.QuotedScalarIndicator {
+						s.prevIndentColumn = tk.Position.Column
+					}
 				}
 				ctx.addToken(token.MappingValue(s.pos()))
 				s.progressColumn(ctx, 1)
@@ -805,6 +811,11 @@ func (s *Scanner) scan(ctx *Context) (pos int) {
 				token, progress := s.scanQuote(ctx, c)
 				ctx.addToken(token)
 				pos += progress
+				// If the non-whitespace character immediately following the quote is ':', the quote should be treated as a map key.
+				// Therefore, do not return and continue processing as a normal map key.
+				if ctx.currentCharWithSkipWhitespace() == ':' {
+					continue
+				}
 				return
 			}
 		case '\r', '\n':