yaml: Support multi document decode/encode

Related to #1087

Author: wader

doc/usage.md

@@ -814,9 +814,12 @@ JSON and jq-flavoured JSON
 Note that `fromjson` and `tojson` use different naming conventions as they originate from jq's standard library.
 
 YAML
-- `from_yaml` Parse YAML into jq value.
+- `from_yaml`/`from_yaml($opts)` Parse YAML into jq value.<br>
+  `$opts` are:
+  - `{multi_document: boolean}` Force multi document mode.<br>
 - `to_yaml`/`to_yaml($opts)` Serialize jq value into YAML. `$opts` are:
   - `{indent: number}` Indent depth.
+  - `{multi_document: boolean}` Force multi document mode.<br>
 
 TOML
 - `from_toml` Parse TOML into jq value.

format/format.go

@@ -408,3 +408,7 @@ type Pg_Heap_In struct {
 type Pg_BTree_In struct {
 	Page int `doc:"First page number in file, default is 0"`
 }
+
+type YAML_In struct {
+	MultiDocument bool `doc:"Force multi document"`
+}

format/yaml/testdata/multi.fqtest

@@ -0,0 +1,44 @@
+$ fq -d yaml . single.yaml
+[]
+$ fq -d yaml . multi.yaml
+[
+  1,
+  2,
+  3
+]
+$ fq -o multi_document=true -d yaml . single.yaml
+[
+  []
+]
+$ fq -o multi_document=true -d yaml . multi.yaml
+[
+  1,
+  2,
+  3
+]
+$ fq -o multi_document=false -d yaml . single.yaml
+[]
+$ fq -o multi_document=false -d yaml . multi.yaml
+[
+  1,
+  2,
+  3
+]
+$ fq -r -o multi_document=true -d yaml '. as $i | to_yaml({multi_document: true}), $i == from_yaml' multi.yaml
+1
+---
+2
+---
+3
+
+true
+$ fq -nr '"[]" | from_yaml({multi_document: true})'
+[
+  []
+]
+$ fq -nr '"[]" | from_yaml({multi_document: false})'
+[]
+$ fq -n '123 | to_yaml({multi_document: true})'
+exitcode: 5
+stderr:
+error: to_yaml cannot be applied to: number (123)

format/yaml/testdata/multi.yaml

@@ -0,0 +1,5 @@
+1
+---
+2
+---
+3

format/yaml/testdata/single.yaml

@@ -0,0 +1 @@
+[]

format/yaml/testdata/trailing.fqtest

@@ -5,4 +5,4 @@ error: error at position 0xc: yaml: line 2: could not find expected ':'
 $ fq -n '`{"a":123}{"b":444}` | from_yaml'
 exitcode: 5
 stderr:
-error: error at position 0x12: trialing data after top-level value
+error: error at position 0x12: trialing data after document

format/yaml/yaml.go

@@ -28,26 +28,46 @@ func init() {
 			ProbeOrder:  format.ProbeOrderTextFuzzy,
 			Groups:      []*decode.Group{format.Probe},
 			DecodeFn:    decodeYAML,
-			Functions:   []string{"_todisplay"},
+			DefaultInArg: format.YAML_In{
+				MultiDocument: false,
+			},
+			Functions: []string{"_todisplay"},
 		})
 	interp.RegisterFS(yamlFS)
 	interp.RegisterFunc1("_to_yaml", toYAML)
 }
 
 func decodeYAML(d *decode.D) any {
+	var yi format.YAML_In
+	d.ArgAs(&yi)
+
 	br := d.RawLen(d.Len())
-	var r any
+
+	var vs []any
 
 	yd := yaml.NewDecoder(bitio.NewIOReader(br))
-	if err := yd.Decode(&r); err != nil {
-		d.Fatalf("%s", err)
-	}
-	if err := yd.Decode(new(any)); !errors.Is(err, io.EOF) {
-		d.Fatalf("trialing data after top-level value")
+	for {
+		var v any
+		err := yd.Decode(&v)
+		if err != nil {
+			if len(vs) == 0 {
+				d.Fatalf("%s", err)
+			} else if errors.Is(err, io.EOF) {
+				break
+			} else {
+				d.Fatalf("trialing data after document")
+			}
+		}
+
+		vs = append(vs, v)
 	}
 
 	var s scalar.Any
-	s.Actual = gojqx.Normalize(r)
+	if !yi.MultiDocument && len(vs) == 1 {
+		s.Actual = gojqx.Normalize(vs[0])
+	} else {
+		s.Actual = gojqx.Normalize(vs)
+	}
 
 	switch s.Actual.(type) {
 	case map[string]any,
@@ -63,18 +83,32 @@ func decodeYAML(d *decode.D) any {
 }
 
 type ToYAMLOpts struct {
-	Indent int `default:"4"` // 4 is default for gopkg.in/yaml.v3
+	Indent        int  `default:"4"` // 4 is default for gopkg.in/yaml.v3
+	MultiDocument bool `default:"false"`
 }
 
 func toYAML(_ *interp.Interp, c any, opts ToYAMLOpts) any {
+	c = gojqx.Normalize(c)
+
+	cs, isArray := c.([]any)
+	if opts.MultiDocument {
+		if !isArray {
+			return gojqx.FuncTypeError{Name: "to_yaml", V: c}
+		}
+	} else {
+		cs = []any{c}
+	}
+
 	b := &bytes.Buffer{}
 	e := yaml.NewEncoder(b)
-	// yaml.SetIndent panics if < 0
-	if opts.Indent >= 0 {
-		e.SetIndent(opts.Indent)
-	}
-	if err := e.Encode(gojqx.Normalize(c)); err != nil {
-		return err
+	for _, c := range cs {
+		// yaml.SetIndent panics if < 0
+		if opts.Indent >= 0 {
+			e.SetIndent(opts.Indent)
+		}
+		if err := e.Encode(gojqx.Normalize(c)); err != nil {
+			return err
+		}
 	}
 
 	return b.String()