From 83d3e55c2408552b08bd69ac9b0b5f6d6e47d9d7 Mon Sep 17 00:00:00 2001
From: Shuhei Kitagawa <shuheiktgw@users.noreply.github.com>
Date: Thu, 12 Dec 2024 10:05:46 +0100
Subject: [PATCH] Fix quotation for timestamps (#515)

Co-authored-by: shuheiktgw <s-kitagawa@mercari.com>
---
 encode_test.go      | 20 +++++++++++---------
 token/token.go      | 33 ++++++++++++++++++++-------------
 token/token_test.go | 11 +++++++++--
 3 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/encode_test.go b/encode_test.go
index 2917c857..06b097ef 100644
--- a/encode_test.go
+++ b/encode_test.go
@@ -273,7 +273,7 @@ func TestEncoder(t *testing.T) {
 			nil,
 		},
 		{
-			"t2: 2018-01-09T10:40:47Z\nt4: 2098-01-09T10:40:47Z\n",
+			"t2: \"2018-01-09T10:40:47Z\"\nt4: \"2098-01-09T10:40:47Z\"\n",
 			map[string]string{
 				"t2": "2018-01-09T10:40:47Z",
 				"t4": "2098-01-09T10:40:47Z",
@@ -717,14 +717,16 @@ func TestEncoder(t *testing.T) {
 		},
 	}
 	for _, test := range tests {
-		var buf bytes.Buffer
-		enc := yaml.NewEncoder(&buf, test.options...)
-		if err := enc.Encode(test.value); err != nil {
-			t.Fatalf("%+v", err)
-		}
-		if test.source != buf.String() {
-			t.Fatalf("expect = [%s], actual = [%s]", test.source, buf.String())
-		}
+		t.Run(test.source, func(t *testing.T) {
+			var buf bytes.Buffer
+			enc := yaml.NewEncoder(&buf, test.options...)
+			if err := enc.Encode(test.value); err != nil {
+				t.Fatalf("%+v", err)
+			}
+			if test.source != buf.String() {
+				t.Fatalf("expect = [%s], actual = [%s]", test.source, buf.String())
+			}
+		})
 	}
 }
 
diff --git a/token/token.go b/token/token.go
index 42e763a5..b880bd28 100644
--- a/token/token.go
+++ b/token/token.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"time"
 )
 
 // Character type for character
@@ -612,20 +613,26 @@ func ToNumber(value string) *NumberValue {
 	}
 }
 
-func looksLikeTimeValue(value string) bool {
-	for i, c := range value {
-		switch c {
-		case ':', '1', '2', '3', '4', '5', '6', '7', '8', '9':
-			continue
-		case '0':
-			if i == 0 {
-				return false
-			}
-			continue
+// This is a subset of the formats permitted by the regular expression
+// defined at http://yaml.org/type/timestamp.html. Note that time.Parse
+// cannot handle: "2001-12-14 21:59:43.10 -5" from the examples.
+var timestampFormats = []string{
+	time.RFC3339Nano,
+	"2006-01-02t15:04:05.999999999Z07:00", // RFC3339Nano with lower-case "t".
+	time.DateTime,
+	time.DateOnly,
+
+	// Not in examples, but to preserve backward compatibility by quoting time values.
+	"15:4",
+}
+
+func isTimestamp(value string) bool {
+	for _, format := range timestampFormats {
+		if _, err := time.Parse(format, value); err == nil {
+			return true
 		}
-		return false
 	}
-	return true
+	return false
 }
 
 // IsNeedQuoted whether need quote for passed string or not
@@ -649,7 +656,7 @@ func IsNeedQuoted(value string) bool {
 	case ':', ' ':
 		return true
 	}
-	if looksLikeTimeValue(value) {
+	if isTimestamp(value) {
 		return true
 	}
 	for i, c := range value {
diff --git a/token/token_test.go b/token/token_test.go
index 27565fbb..e5e5ba8f 100644
--- a/token/token_test.go
+++ b/token/token_test.go
@@ -82,6 +82,10 @@ func TestIsNeedQuoted(t *testing.T) {
 		"true",
 		"1.234",
 		"1:1",
+		"2001-12-15T02:59:43.1Z",
+		"2001-12-14t21:59:43.10-05:00",
+		"2001-12-15 2:59:43.10",
+		"2002-12-14",
 		"hoge # comment",
 		"\\0",
 		"#a b",
@@ -128,15 +132,18 @@ func TestIsNeedQuoted(t *testing.T) {
 	}
 	for i, test := range needQuotedTests {
 		if !token.IsNeedQuoted(test) {
-			t.Fatalf("%d: failed to quoted judge for %s", i, test)
+			t.Errorf("%d: failed to quoted judge for %s", i, test)
 		}
 	}
 	notNeedQuotedTests := []string{
 		"Hello World",
+		// time.Parse cannot handle: "2001-12-14 21:59:43.10 -5" from the examples.
+		// https://yaml.org/type/timestamp.html
+		"2001-12-14 21:59:43.10 -5",
 	}
 	for i, test := range notNeedQuotedTests {
 		if token.IsNeedQuoted(test) {
-			t.Fatalf("%d: failed to quoted judge for %s", i, test)
+			t.Errorf("%d: failed to quoted judge for %s", i, test)
 		}
 	}
 }