Implement Encoder and Decoder

yaml.go

@@ -161,7 +161,7 @@ func yamlToJSONTarget(yamlBytes []byte, jsonTarget *reflect.Value, unmarshalFn f
 	// 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
-	// incompatibilties happen along the way.
+	// incompatibilities happen along the way.
 	jsonObj, err := convertToJSONableObject(yamlObj, jsonTarget)
 	if err != nil {
 		return nil, fmt.Errorf("error converting YAML to JSON: %w", err)
@@ -414,3 +414,113 @@ func jsonToYAMLValue(j interface{}) interface{} {
 	}
 	return j
 }
+
+// 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
+}
+
+// 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
+}
+
+// SetStrict sets whether strict decoding behaviour is enabled when decoding items in the data (see UnmarshalStrict).
+// By default, decoding is not strict. It also adds the DisallowUnknownFields option to json decoder.
+func (dec *Decoder) SetStrict() {
+	dec.decoder.SetStrict(true)
+	dec.opts = append(dec.opts, DisallowUnknownFields)
+}

yaml_test.go

@@ -17,16 +17,20 @@ limitations under the License.
 package yaml
 
 import (
+	"bytes"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"io"
 	"math"
 	"reflect"
 	"sort"
 	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/davecgh/go-spew/spew"
-	yaml "gopkg.in/yaml.v2"
+	"gopkg.in/yaml.v2"
 )
 
 /* Test helper functions */
@@ -59,6 +63,16 @@ var (
 	funcUnmarshalStrict testUnmarshalFunc = func(yamlBytes []byte, obj interface{}) error {
 		return UnmarshalStrict(yamlBytes, obj)
 	}
+
+	funcDecode testUnmarshalFunc = func(yamlBytes []byte, obj interface{}) error {
+		return NewDecoder(bytes.NewReader(yamlBytes)).Decode(obj)
+	}
+
+	funcDecodeStrict testUnmarshalFunc = func(yamlBytes []byte, obj interface{}) error {
+		decoder := NewDecoder(bytes.NewReader(yamlBytes))
+		decoder.SetStrict()
+		return decoder.Decode(obj)
+	}
 )
 
 func testUnmarshal(t *testing.T, f testUnmarshalFunc, tests map[string]unmarshalTestCase) {
@@ -168,6 +182,43 @@ func testYAMLToJSON(t *testing.T, f testYAMLToJSONFunc, tests map[string]yamlToJ
 	}
 }
 
+type testMarshalFunc = func(obj interface{}) ([]byte, error)
+
+var (
+	funcMarshal testMarshalFunc = func(obj interface{}) ([]byte, error) {
+		return Marshal(obj)
+	}
+	funcEncode testMarshalFunc = func(obj interface{}) ([]byte, error) {
+		var buf bytes.Buffer
+		encoder := NewEncoder(&buf)
+		if err := encoder.Encode(obj); err != nil {
+			return nil, err
+		}
+		return buf.Bytes(), nil
+	}
+)
+
+type marshalTestCase struct {
+	obj     interface{}
+	encoded []byte
+}
+
+func testMarshal(t *testing.T, f testMarshalFunc, tests map[string]marshalTestCase) {
+	for testName, test := range tests {
+		t.Run(testName, func(t *testing.T) {
+			y, err := f(test.obj)
+			if err != nil {
+				t.Errorf("error marshaling YAML: %v", err)
+			}
+
+			if !reflect.DeepEqual(y, test.encoded) {
+				t.Errorf("marshal YAML was unsuccessful, expected: %#v, got: %#v",
+					string(test.encoded), string(y))
+			}
+		})
+	}
+}
+
 /* Start tests */
 
 type MarshalTest struct {
@@ -180,18 +231,20 @@ type MarshalTest struct {
 
 func TestMarshal(t *testing.T) {
 	f32String := strconv.FormatFloat(math.MaxFloat32, 'g', -1, 32)
-	s := MarshalTest{"a", math.MaxInt64, math.MaxFloat32}
-	e := []byte(fmt.Sprintf("A: a\nB: %d\nC: %s\n", math.MaxInt64, f32String))
-
-	y, err := Marshal(s)
-	if err != nil {
-		t.Errorf("error marshaling YAML: %v", err)
+	tests := map[string]marshalTestCase{
+		"max": {
+			obj:     MarshalTest{"a", math.MaxInt64, math.MaxFloat32},
+			encoded: []byte(fmt.Sprintf("A: a\nB: %d\nC: %s\n", math.MaxInt64, f32String)),
+		},
 	}
 
-	if !reflect.DeepEqual(y, e) {
-		t.Errorf("marshal YAML was unsuccessful, expected: %#v, got: %#v",
-			string(e), string(y))
-	}
+	t.Run("Marshal", func(t *testing.T) {
+		testMarshal(t, funcMarshal, tests)
+	})
+
+	t.Run("Encode", func(t *testing.T) {
+		testMarshal(t, funcEncode, tests)
+	})
 }
 
 type UnmarshalUntaggedStruct struct {
@@ -550,6 +603,14 @@ func TestUnmarshal(t *testing.T) {
 	t.Run("UnmarshalStrict", func(t *testing.T) {
 		testUnmarshal(t, funcUnmarshalStrict, tests)
 	})
+
+	t.Run("Decode", func(t *testing.T) {
+		testUnmarshal(t, funcDecode, tests)
+	})
+
+	t.Run("DecodeStrict", func(t *testing.T) {
+		testUnmarshal(t, funcDecodeStrict, tests)
+	})
 }
 
 func TestUnmarshalStrictFails(t *testing.T) {
@@ -618,6 +679,19 @@ func TestUnmarshalStrictFails(t *testing.T) {
 		}
 		testUnmarshal(t, funcUnmarshalStrict, failTests)
 	})
+
+	t.Run("Decode", func(t *testing.T) {
+		testUnmarshal(t, funcDecode, tests)
+	})
+
+	t.Run("DecodeStrict", func(t *testing.T) {
+		failTests := map[string]unmarshalTestCase{}
+		for name, test := range tests {
+			test.err = fatalErrorsType
+			failTests[name] = test
+		}
+		testUnmarshal(t, funcDecodeStrict, failTests)
+	})
 }
 
 func TestYAMLToJSON(t *testing.T) {
@@ -890,3 +964,64 @@ func sortMapSlicesInPlace(x interface{}) {
 		})
 	}
 }
+
+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)
+
+	}
+}