Implement Decoder and Encoder

yaml.go

@@ -307,3 +307,125 @@ func convertToJSONableObject(yamlObj interface{}, jsonTarget *reflect.Value) (in
 		return yamlObj, nil
 	}
 }
+
+// An Encoder writes YAML values to an output stream.
+type Encoder struct {
+	encoder *yaml.Encoder
+}
+
+// NewEncoder returns a new encoder that writes to w. The Encoder should be closed after use to flush all data to w.
+func NewEncoder(w io.Writer) *Encoder {
+	return &Encoder{
+		encoder: yaml.NewEncoder(w),
+	}
+}
+
+// Encode writes the YAML encoding of obj to the stream.
+// If multiple items are encoded to the stream, the second and subsequent document will be preceded with a "---" document separator,
+// but the first will not.
+//
+// See the documentation for Marshal for details about the conversion of Go values to YAML.
+func (e *Encoder) Encode(obj interface{}) error {
+	var buf bytes.Buffer
+	// Convert an object to the JSON.
+	if err := json.NewEncoder(&buf).Encode(obj); err != nil {
+		return fmt.Errorf("error encode into JSON: %w", err)
+	}
+
+	// Convert the JSON to an object.
+	var j interface{}
+	// We are using yaml.Decoder.Decode here (instead of json.Decoder.Decode) because the
+	// Go JSON library doesn't try to pick the right number type (int, float,
+	// etc.) when unmarshalling to interface{}, it just picks float64
+	// universally. go-yaml does go through the effort of picking the right
+	// number type, so we can preserve number type throughout this process.
+	if err := yaml.NewDecoder(&buf).Decode(&j); err != nil {
+		return fmt.Errorf("error decode from JSON: %w", err)
+	}
+	// Marshal this object into YAML.
+	if err := e.encoder.Encode(j); err != nil {
+		return fmt.Errorf("error encode into YAML: %w", err)
+	}
+
+	return nil
+}
+
+// SetIndent changes the used indentation used when encoding.
+func (e *Encoder) SetIndent(spaces int) {
+	e.encoder.SetIndent(spaces)
+}
+
+// Close closes the encoder by writing any remaining data. It does not write a stream terminating string "...".
+func (e *Encoder) Close() (err error) {
+	if err := e.encoder.Close(); err != nil {
+		return fmt.Errorf("error closing encoder: %w", err)
+	}
+	return nil
+}
+
+// A Decoder reads and decodes YAML values from an input stream.
+type Decoder struct {
+	opts    []JSONOpt
+	decoder *yaml.Decoder
+}
+
+// NewDecoder returns a new decoder that reads from r.
+//
+// The decoder introduces its own buffering and may read data from r beyond the YAML values requested.
+//
+// Options for the standard library json.Decoder can be optionally specified, e.g. to decode untyped numbers into json.Number instead of float64,
+// or to disallow unknown fields (but for that purpose, see also UnmarshalStrict)
+func NewDecoder(r io.Reader, opts ...JSONOpt) *Decoder {
+	return &Decoder{
+		opts:    opts,
+		decoder: yaml.NewDecoder(r),
+	}
+}
+
+// Decode reads the next YAML-encoded value from its input and stores it in the value pointed to by obj.
+//
+// See the documentation for Unmarshal for details about the conversion of YAML into a Go value.
+func (dec *Decoder) Decode(obj interface{}) error {
+	jsonTarget := reflect.ValueOf(obj)
+
+	// Convert the YAML to an object.
+	var yamlObj interface{}
+	if err := dec.decoder.Decode(&yamlObj); err != nil {
+		return fmt.Errorf("error converting YAML to JSON: %w", err)
+	}
+
+	// YAML objects are not completely compatible with JSON objects (e.g. you
+	// can have non-string keys in YAML). So, convert the YAML-compatible object
+	// to a JSON-compatible object, failing with an error if irrecoverable
+	// incompatibilities happen along the way.
+	jsonObj, err := convertToJSONableObject(yamlObj, &jsonTarget)
+	if err != nil {
+		return fmt.Errorf("error converting YAML to JSON: %w", err)
+	}
+
+	// Convert this object to JSON
+	var buf bytes.Buffer
+	if err := json.NewEncoder(&buf).Encode(jsonObj); err != nil {
+		return fmt.Errorf("error converting YAML to JSON: %w", err)
+	}
+
+	if err := jsonUnmarshal(&buf, obj, dec.opts...); err != nil {
+		return fmt.Errorf("error unmarshaling JSON: %w", err)
+	}
+
+	return nil
+}
+
+// KnownFields ensures that the keys in decoded mappings to
+// exist as fields in the struct being decoded into.
+func (dec *Decoder) KnownFields() {
+	dec.decoder.KnownFields(true)
+	dec.opts = append(dec.opts, DisallowUnknownFields)
+}
+
+// DisallowUnknownFields configures the JSON decoder to error out if unknown
+// fields come along, instead of dropping them by default.
+func DisallowUnknownFields(d *json.Decoder) *json.Decoder {
+	d.DisallowUnknownFields()
+	return d
+}

yaml_test.go

@@ -1,7 +1,10 @@
 package yaml
 
 import (
+	"bytes"
+	"errors"
 	"fmt"
+	"io"
 	"math"
 	"reflect"
 	"runtime"
@@ -23,15 +26,29 @@ func TestMarshal(t *testing.T) {
 	s := MarshalTest{"a", math.MaxInt64, math.MaxFloat32, math.MaxFloat64}
 	e := []byte(fmt.Sprintf("A: a\nB: %d\nC: %s\nD: %s\n", int64(math.MaxInt64), f32String, f64String))
 
-	y, err := Marshal(s)
-	if err != nil {
-		t.Errorf("error marshaling YAML: %v", err)
-	}
+	t.Run("Marshal", func(t *testing.T) {
+		y, err := Marshal(s)
+		if err != nil {
+			t.Errorf("error marshaling YAML: %v", err)
+		}
+
+		if !reflect.DeepEqual(y, e) {
+			t.Errorf("marshal YAML was unsuccessful, expected: %#v, got: %#v",
+				string(e), string(y))
+		}
+	})
+	t.Run("Encode", func(t *testing.T) {
+		var buf bytes.Buffer
+		if err := NewEncoder(&buf).Encode(s); err != nil {
+			t.Errorf("error encoding YAML: %v", err)
+		}
+
+		if y := buf.Bytes(); !reflect.DeepEqual(y, e) {
+			t.Errorf("encode YAML was unsuccessful, expected: %#v, got: %#v",
+				string(e), string(y))
+		}
+	})
 
-	if !reflect.DeepEqual(y, e) {
-		t.Errorf("marshal YAML was unsuccessful, expected: %#v, got: %#v",
-			string(e), string(y))
-	}
 }
 
 type UnmarshalString struct {
@@ -160,16 +177,26 @@ func prettyFunctionName(opts []JSONOpt) []string {
 }
 
 func unmarshalEqual(t *testing.T, y []byte, s, e interface{}, opts ...JSONOpt) { //nolint:unparam
-	t.Helper()
-	err := Unmarshal(y, s, opts...)
-	if err != nil {
-		t.Errorf("Unmarshal(%#q, s, %v) = %v", string(y), prettyFunctionName(opts), err)
-		return
-	}
+	t.Run("Unmarshal", func(t *testing.T) {
+		if err := Unmarshal(y, s, opts...); err != nil {
+			t.Errorf("Unmarshal(%#q, s, %v) = %v", string(y), prettyFunctionName(opts), err)
+			return
+		}
 
-	if !reflect.DeepEqual(s, e) {
-		t.Errorf("Unmarshal(%#q, s, %v) = %+#v; want %+#v", string(y), prettyFunctionName(opts), s, e)
-	}
+		if !reflect.DeepEqual(s, e) {
+			t.Errorf("Unmarshal(%#q, s, %v) = %+#v; want %+#v", string(y), prettyFunctionName(opts), s, e)
+		}
+	})
+	t.Run("Decode", func(t *testing.T) {
+		if err := NewDecoder(bytes.NewReader(y), opts...).Decode(s); err != nil {
+			t.Errorf("Decode(%#q, s, %v) = %v", string(y), prettyFunctionName(opts), err)
+			return
+		}
+
+		if !reflect.DeepEqual(s, e) {
+			t.Errorf("Decode(%#q, s, %v) = %+#v; want %+#v", string(y), prettyFunctionName(opts), s, e)
+		}
+	})
 }
 
 // TestUnmarshalErrors tests that we return an error on ambiguous YAML.
@@ -396,3 +423,64 @@ func TestYAMLToJSONDuplicateFields(t *testing.T) {
 		t.Error("expected YAMLtoJSON to fail on duplicate field names")
 	}
 }
+
+type MultiDoc struct {
+	Test int `json:"test"`
+}
+
+func TestMultiDocDecode(t *testing.T) {
+	data := `---
+test: 1
+---
+test: 2
+---
+test: 3
+`
+	decoder := NewDecoder(strings.NewReader(data))
+
+	for i := 1; i < 4; i++ {
+		var obj MultiDoc
+		if err := decoder.Decode(&obj); err != nil {
+			t.Errorf("decode #%d failed: %s", i, err)
+		}
+		if obj.Test != i {
+			t.Errorf("decoded #%d has incorrect value %#v", i, obj)
+		}
+	}
+	var obj MultiDoc
+	if err := decoder.Decode(&obj); !errors.Is(err, io.EOF) {
+		t.Errorf("decode should return EOF but got: %s", err)
+	}
+}
+
+func TestMultiDocEncode(t *testing.T) {
+	docs := []MultiDoc{
+		{
+			Test: 1,
+		},
+		{
+			Test: 2,
+		},
+		{
+			Test: 3,
+		},
+	}
+	expected := `test: 1
+---
+test: 2
+---
+test: 3
+`
+
+	var buf bytes.Buffer
+	encoder := NewEncoder(&buf)
+	for _, obj := range docs {
+		if err := encoder.Encode(obj); err != nil {
+			t.Errorf("encode object %#v failed: %s", obj, err)
+		}
+	}
+	if encoded := buf.String(); encoded != expected {
+		t.Errorf("expected %s, but got %s", expected, encoded)
+
+	}
+}