From 3ac9c8bc4174656831ec2e7c4e6fcfbf98ffa5e3 Mon Sep 17 00:00:00 2001
From: Tom Wright <contact@tomwright.me>
Date: Sun, 27 Oct 2024 16:47:22 +0000
Subject: [PATCH] Reimplement yaml parser

---
 model/value.go            |   9 +-
 parsing/yaml/yaml.go      | 242 ++++++++++++++++++++++++++++++++++++--
 parsing/yaml/yaml_test.go | 186 +++++++++++++++++++++++++++++
 3 files changed, 427 insertions(+), 10 deletions(-)
 create mode 100644 parsing/yaml/yaml_test.go

diff --git a/model/value.go b/model/value.go
index 623edc9..40cb873 100644
--- a/model/value.go
+++ b/model/value.go
@@ -47,7 +47,8 @@ func NewValue(v any) *Value {
 		return val
 	case reflect.Value:
 		return &Value{
-			Value: val,
+			Value:    val,
+			Metadata: make(map[string]any),
 		}
 	case nil:
 		return NewNullValue()
@@ -57,13 +58,17 @@ func NewValue(v any) *Value {
 			res.Elem().Set(reflect.ValueOf(v))
 		}
 		return &Value{
-			Value: res,
+			Value:    res,
+			Metadata: make(map[string]any),
 		}
 	}
 }
 
 // Interface returns the value as an interface.
 func (v *Value) Interface() any {
+	if v.IsNull() {
+		return nil
+	}
 	return v.Value.Interface()
 }
 
diff --git a/parsing/yaml/yaml.go b/parsing/yaml/yaml.go
index 8f9e433..ce539d8 100644
--- a/parsing/yaml/yaml.go
+++ b/parsing/yaml/yaml.go
@@ -2,6 +2,9 @@ package yaml
 
 import (
 	"bytes"
+	"fmt"
+	"io"
+	"strconv"
 
 	"github.com/tomwright/dasel/v3/model"
 	"github.com/tomwright/dasel/v3/parsing"
@@ -32,21 +35,244 @@ type yamlReader struct{}
 // Read reads a value from a byte slice.
 func (j *yamlReader) Read(data []byte) (*model.Value, error) {
 	d := yaml.NewDecoder(bytes.NewReader(data))
-	var unmarshalled any
-	if err := d.Decode(&unmarshalled); err != nil {
-		return nil, err
+	res := make([]*yamlValue, 0)
+	for {
+		unmarshalled := &yamlValue{}
+		if err := d.Decode(&unmarshalled); err != nil {
+			if err == io.EOF {
+				break
+			}
+			return nil, err
+		}
+		res = append(res, unmarshalled)
+	}
+
+	switch len(res) {
+	case 0:
+		return model.NewNullValue(), nil
+	case 1:
+		return res[0].value, nil
+	default:
+		slice := model.NewSliceValue()
+		slice.MarkAsBranch()
+		for _, v := range res {
+			if err := slice.Append(v.value); err != nil {
+				return nil, err
+			}
+		}
+		return slice, nil
 	}
-	return model.NewValue(&unmarshalled), nil
 }
 
 type yamlWriter struct{}
 
 // Write writes a value to a byte slice.
 func (j *yamlWriter) Write(value *model.Value) ([]byte, error) {
-	buf := new(bytes.Buffer)
-	e := yaml.NewEncoder(buf)
-	if err := e.Encode(value.Interface()); err != nil {
+	if value.IsBranch() {
+		res := make([]byte, 0)
+		sliceLen, err := value.SliceLen()
+		if err != nil {
+			return nil, err
+		}
+		if err := value.RangeSlice(func(i int, val *model.Value) error {
+			yv := &yamlValue{value: val}
+			marshalled, err := yaml.Marshal(yv)
+			if err != nil {
+				return err
+			}
+			res = append(res, marshalled...)
+			if i < sliceLen-1 {
+				res = append(res, []byte("---\n")...)
+			}
+			return nil
+		}); err != nil {
+			return nil, err
+		}
+		return res, nil
+	}
+
+	yv := &yamlValue{value: value}
+	res, err := yv.ToNode()
+	if err != nil {
+		return nil, err
+	}
+	return yaml.Marshal(res)
+}
+
+type yamlValue struct {
+	node  *yaml.Node
+	value *model.Value
+}
+
+func (yv *yamlValue) UnmarshalYAML(value *yaml.Node) error {
+	yv.node = value
+	switch value.Kind {
+	case yaml.ScalarNode:
+		switch value.Tag {
+		case "!!bool":
+			yv.value = model.NewBoolValue(value.Value == "true")
+		case "!!int":
+			i, err := strconv.Atoi(value.Value)
+			if err != nil {
+				return err
+			}
+			yv.value = model.NewIntValue(int64(i))
+		case "!!float":
+			f, err := strconv.ParseFloat(value.Value, 64)
+			if err != nil {
+				return err
+			}
+			yv.value = model.NewFloatValue(f)
+		default:
+			yv.value = model.NewStringValue(value.Value)
+		}
+	case yaml.DocumentNode:
+		yv.value = model.NewNullValue()
+	case yaml.SequenceNode:
+		res := model.NewSliceValue()
+		for _, item := range value.Content {
+			newItem := &yamlValue{}
+			if err := newItem.UnmarshalYAML(item); err != nil {
+				return err
+			}
+			if err := res.Append(newItem.value); err != nil {
+				return err
+			}
+		}
+		yv.value = res
+	case yaml.MappingNode:
+		res := model.NewMapValue()
+		for i := 0; i < len(value.Content); i += 2 {
+			key := value.Content[i]
+			val := value.Content[i+1]
+
+			newKey := &yamlValue{}
+			if err := newKey.UnmarshalYAML(key); err != nil {
+				return err
+			}
+
+			newVal := &yamlValue{}
+			if err := newVal.UnmarshalYAML(val); err != nil {
+				return err
+			}
+
+			keyStr, err := newKey.value.StringValue()
+			if err != nil {
+				return fmt.Errorf("keys are expected to be strings: %w", err)
+			}
+
+			if err := res.SetMapKey(keyStr, newVal.value); err != nil {
+				return err
+			}
+		}
+		yv.value = res
+	case yaml.AliasNode:
+		newVal := &yamlValue{}
+		if err := newVal.UnmarshalYAML(value.Alias); err != nil {
+			return err
+		}
+		yv.value = newVal.value
+		yv.value.Metadata["yaml-alias"] = value.Value
+	}
+	return nil
+}
+
+func (yv *yamlValue) ToNode() (*yaml.Node, error) {
+	res := &yaml.Node{}
+
+	yamlAlias, ok := yv.value.Metadata["yaml-alias"].(string)
+	if ok {
+		//res.Kind = yaml.ScalarNode
+		res.Kind = yaml.AliasNode
+		res.Value = yamlAlias
+		//res.Alias = &yaml.Node{
+		//	Kind:  yaml.ScalarNode,
+		//	Value: yamlAlias,
+		//}
+		return res, nil
+	}
+
+	switch yv.value.Type() {
+	case model.TypeString:
+		v, err := yv.value.StringValue()
+		if err != nil {
+			return nil, err
+		}
+		res.Kind = yaml.ScalarNode
+		res.Value = v
+		res.Tag = "!!str"
+	case model.TypeBool:
+		v, err := yv.value.BoolValue()
+		if err != nil {
+			return nil, err
+		}
+		res.Kind = yaml.ScalarNode
+		res.Value = fmt.Sprintf("%t", v)
+		res.Tag = "!!bool"
+	case model.TypeInt:
+		v, err := yv.value.IntValue()
+		if err != nil {
+			return nil, err
+		}
+		res.Kind = yaml.ScalarNode
+		res.Value = fmt.Sprintf("%d", v)
+		res.Tag = "!!int"
+	case model.TypeFloat:
+		v, err := yv.value.FloatValue()
+		if err != nil {
+			return nil, err
+		}
+		res.Kind = yaml.ScalarNode
+		res.Value = fmt.Sprintf("%g", v)
+		res.Tag = "!!float"
+	case model.TypeMap:
+		res.Kind = yaml.MappingNode
+		if err := yv.value.RangeMap(func(key string, val *model.Value) error {
+			keyNode := &yamlValue{value: model.NewStringValue(key)}
+			valNode := &yamlValue{value: val}
+
+			marshalledKey, err := keyNode.ToNode()
+			if err != nil {
+				return err
+			}
+			marshalledVal, err := valNode.ToNode()
+			if err != nil {
+				return err
+			}
+
+			res.Content = append(res.Content, marshalledKey)
+			res.Content = append(res.Content, marshalledVal)
+
+			return nil
+		}); err != nil {
+			return nil, err
+		}
+	case model.TypeSlice:
+		res.Kind = yaml.SequenceNode
+		if err := yv.value.RangeSlice(func(i int, val *model.Value) error {
+			valNode := &yamlValue{value: val}
+			marshalledVal, err := valNode.ToNode()
+			if err != nil {
+				return err
+			}
+			res.Content = append(res.Content, marshalledVal)
+			return nil
+		}); err != nil {
+			return nil, err
+		}
+	case model.TypeNull:
+		res.Kind = yaml.DocumentNode
+	case model.TypeUnknown:
+		return nil, fmt.Errorf("unknown type: %s", yv.value.Type())
+	}
+
+	return res, nil
+}
+
+func (yv *yamlValue) MarshalYAML() (any, error) {
+	res, err := yv.ToNode()
+	if err != nil {
 		return nil, err
 	}
-	return buf.Bytes(), nil
+	return res, nil
 }
diff --git a/parsing/yaml/yaml_test.go b/parsing/yaml/yaml_test.go
new file mode 100644
index 0000000..264e8ab
--- /dev/null
+++ b/parsing/yaml/yaml_test.go
@@ -0,0 +1,186 @@
+package yaml_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/tomwright/dasel/v3/model"
+	"github.com/tomwright/dasel/v3/parsing"
+	"github.com/tomwright/dasel/v3/parsing/yaml"
+)
+
+type testCase struct {
+	in     string
+	assert func(t *testing.T, res *model.Value)
+}
+
+func (tc testCase) run(t *testing.T) {
+	r, err := yaml.YAML.NewReader()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	res, err := r.Read([]byte(tc.in))
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	tc.assert(t, res)
+}
+
+type rwTestCase struct {
+	in  string
+	out string
+}
+
+func (tc rwTestCase) run(t *testing.T) {
+	if tc.out == "" {
+		tc.out = tc.in
+	}
+	r, err := yaml.YAML.NewReader()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	w, err := yaml.YAML.NewWriter(parsing.WriterOptions{})
+	res, err := r.Read([]byte(tc.in))
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	out, err := w.Write(res)
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	if !bytes.Equal([]byte(tc.out), out) {
+		t.Errorf("unexpected output: %s", cmp.Diff(tc.out, string(out)))
+	}
+}
+
+func TestYamlValue_UnmarshalYAML(t *testing.T) {
+	t.Run("simple key value", testCase{
+		in: `name: Tom`,
+		assert: func(t *testing.T, res *model.Value) {
+			got, err := res.GetMapKey("name")
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+
+			gotStr, err := got.StringValue()
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+
+			if gotStr != "Tom" {
+				t.Errorf("unexpected value: %s", gotStr)
+			}
+		},
+	}.run)
+
+	t.Run("multi document", testCase{
+		in: `name: Tom
+---
+name: Jerry`,
+		assert: func(t *testing.T, res *model.Value) {
+			a, err := res.GetSliceIndex(0)
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			got, err := a.GetMapKey("name")
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			gotStr, err := got.StringValue()
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			if gotStr != "Tom" {
+				t.Errorf("unexpected value: %s", gotStr)
+			}
+
+			b, err := res.GetSliceIndex(1)
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			got, err = b.GetMapKey("name")
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			gotStr, err = got.StringValue()
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			if gotStr != "Jerry" {
+				t.Errorf("unexpected value: %s", gotStr)
+			}
+		},
+	}.run)
+
+	t.Run("multi document", testCase{
+		in: `name: Tom
+---
+name: Jerry`,
+		assert: func(t *testing.T, res *model.Value) {
+			a, err := res.GetSliceIndex(0)
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			got, err := a.GetMapKey("name")
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			gotStr, err := got.StringValue()
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			if gotStr != "Tom" {
+				t.Errorf("unexpected value: %s", gotStr)
+			}
+
+			b, err := res.GetSliceIndex(1)
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			got, err = b.GetMapKey("name")
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			gotStr, err = got.StringValue()
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			if gotStr != "Jerry" {
+				t.Errorf("unexpected value: %s", gotStr)
+			}
+		},
+	}.run)
+
+	t.Run("multi document", rwTestCase{
+		in: `name: Tom
+---
+name: Jerry
+`,
+	}.run)
+
+	t.Run("generic", rwTestCase{
+		in: `str: foo
+int: 1
+float: 1.1
+bool: true
+map:
+    key: value
+list:
+    - item1
+    - item2
+`,
+	}.run)
+
+	// This test is technically wrong because we're only supporting the alias on read and not write.
+	t.Run("alias", rwTestCase{
+		in: `name: &name Tom
+name2: *name
+`,
+		out: `name: Tom
+name2: Tom
+`,
+	}.run)
+}