Fixed handling of --- #1890, #1896

mikefarah · Dec 12, 2023 · 7f26d91 · 7f26d91
1 parent fa322ae
commit 7f26d91
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git a/pkg/yqlib/decoder_yaml.go b/pkg/yqlib/decoder_yaml.go
@@ -37,7 +37,7 @@ func (dec *yamlDecoder) processReadStream(reader *bufio.Reader) (io.Reader, stri
 	var yamlDirectiveLineRegEx = regexp.MustCompile(`^\s*%YA`)
 	var sb strings.Builder
 	for {
-		peekBytes, err := reader.Peek(3)
+		peekBytes, err := reader.Peek(4)
 		if errors.Is(err, io.EOF) {
 			// EOF are handled else where..
 			return reader, sb.String(), nil
@@ -51,7 +51,15 @@ func (dec *yamlDecoder) processReadStream(reader *bufio.Reader) (io.Reader, stri
 			} else if err != nil {
 				return reader, sb.String(), err
 			}
-		} else if string(peekBytes) == "---" {
+		} else if string(peekBytes) == "--- " {
+			_, err := reader.ReadString(' ')
+			sb.WriteString("$yqDocSeparator$\n")
+			if errors.Is(err, io.EOF) {
+				return reader, sb.String(), nil
+			} else if err != nil {
+				return reader, sb.String(), err
+			}
+		} else if string(peekBytes) == "---\n" {
 			_, err := reader.ReadString('\n')
 			sb.WriteString("$yqDocSeparator$\n")
 			if errors.Is(err, io.EOF) {

diff --git a/pkg/yqlib/yaml_test.go b/pkg/yqlib/yaml_test.go
@@ -7,6 +7,18 @@ import (
 )
 
 var yamlFormatScenarios = []formatScenario{
+	{
+		description: "scalar with doc separator",
+		skipDoc:     true,
+		input:       "--- cat",
+		expected:    "---\ncat\n",
+	},
+	{
+		description: "scalar with doc separator",
+		skipDoc:     true,
+		input:       "---cat",
+		expected:    "---cat\n",
+	},
 	{
 		description: "basic - null",
 		skipDoc:     true,