Skip to content

Commit

Permalink
Support stream decoding
Browse files Browse the repository at this point in the history
  • Loading branch information
goccy committed Jun 1, 2020
1 parent 70a3cab commit 68532a4
Show file tree
Hide file tree
Showing 3 changed files with 107 additions and 16 deletions.
77 changes: 61 additions & 16 deletions decode.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ type Decoder struct {
validator StructValidator
disallowUnknownField bool
useOrderedMap bool
parsedFile *ast.File
streamIndex int
}

// NewDecoder returns a new decoder that reads from r.
Expand Down Expand Up @@ -1131,50 +1133,93 @@ func (d *Decoder) resolveReference() error {
}

// assign new anchor definition to anchorMap
if _, err := d.decode(bytes); err != nil {
if _, err := d.parse(bytes); err != nil {
return errors.Wrapf(err, "failed to decode")
}
}
d.isResolvedReference = true
return nil
}

func (d *Decoder) decode(bytes []byte) (ast.Node, error) {
func (d *Decoder) parse(bytes []byte) (*ast.File, error) {
f, err := parser.ParseBytes(bytes, 0)
if err != nil {
return nil, errors.Wrapf(err, "failed to parse yaml")
}
return d.fileToNode(f), nil
normalizedFile := &ast.File{}
for _, doc := range f.Docs {
// try to decode ast.Node to value and map anchor value to anchorMap
if v := d.nodeToValue(doc.Body); v != nil {
normalizedFile.Docs = append(normalizedFile.Docs, doc)
}
}
return normalizedFile, nil
}

// Decode reads the next YAML-encoded value from its input
// and stores it in the value pointed to by v.
//
// See the documentation for Unmarshal for details about the
// conversion of YAML into a Go value.
func (d *Decoder) Decode(v interface{}) error {
func (d *Decoder) isInitialized() bool {
return d.parsedFile != nil
}

func (d *Decoder) decodeInit() error {
if !d.isResolvedReference {
if err := d.resolveReference(); err != nil {
return errors.Wrapf(err, "failed to resolve reference")
}
}
rv := reflect.ValueOf(v)
if rv.Type().Kind() != reflect.Ptr {
return errors.ErrDecodeRequiredPointerType
}
var buf bytes.Buffer
if _, err := io.Copy(&buf, d.reader); err != nil {
return errors.Wrapf(err, "failed to copy from reader")
}
node, err := d.decode(buf.Bytes())
file, err := d.parse(buf.Bytes())
if err != nil {
return errors.Wrapf(err, "failed to decode")
}
if node == nil {
d.parsedFile = file
return nil
}

func (d *Decoder) decode(v reflect.Value) error {
if len(d.parsedFile.Docs) <= d.streamIndex {
return io.EOF
}
body := d.parsedFile.Docs[d.streamIndex].Body
if body == nil {
return nil
}
if err := d.decodeValue(rv.Elem(), node); err != nil {
if err := d.decodeValue(v.Elem(), body); err != nil {
return errors.Wrapf(err, "failed to decode value")
}
d.streamIndex++
return nil
}

// Decode reads the next YAML-encoded value from its input
// and stores it in the value pointed to by v.
//
// See the documentation for Unmarshal for details about the
// conversion of YAML into a Go value.
func (d *Decoder) Decode(v interface{}) error {
rv := reflect.ValueOf(v)
if rv.Type().Kind() != reflect.Ptr {
return errors.ErrDecodeRequiredPointerType
}
if d.isInitialized() {
if err := d.decode(rv); err != nil {
if err == io.EOF {
return err
}
return errors.Wrapf(err, "failed to decode")
}
return nil
}
if err := d.decodeInit(); err != nil {
return errors.Wrapf(err, "failed to decodInit")
}
if err := d.decode(rv); err != nil {
if err == io.EOF {
return err
}
return errors.Wrapf(err, "failed to decode")
}
return nil
}
42 changes: 42 additions & 0 deletions decode_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package yaml_test
import (
"bytes"
"fmt"
"io"
"log"
"math"
"reflect"
Expand Down Expand Up @@ -921,6 +922,9 @@ func TestDecoder(t *testing.T) {
typ := reflect.ValueOf(test.value).Type()
value := reflect.New(typ)
if err := dec.Decode(value.Interface()); err != nil {
if err == io.EOF {
continue
}
t.Fatalf("%s: %+v", test.source, err)
}
actual := fmt.Sprintf("%+v", value.Elem().Interface())
Expand Down Expand Up @@ -1728,3 +1732,41 @@ j: k
t.Fatalf("expected:[%s] actual:[%s]", string(yml), "\n"+string(bytes))
}
}

func TestDecoder_Stream(t *testing.T) {
yml := `
---
a: b
c: d
---
e: f
g: h
---
i: j
k: l
`
dec := yaml.NewDecoder(strings.NewReader(yml))
values := []map[string]string{}
for {
var v map[string]string
if err := dec.Decode(&v); err != nil {
if err == io.EOF {
break
}
t.Fatalf("%+v", err)
}
values = append(values, v)
}
if len(values) != 3 {
t.Fatal("failed to stream decoding")
}
if values[0]["a"] != "b" {
t.Fatal("failed to stream decoding")
}
if values[1]["e"] != "f" {
t.Fatal("failed to stream decoding")
}
if values[2]["i"] != "j" {
t.Fatal("failed to stream decoding")
}
}
4 changes: 4 additions & 0 deletions yaml.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package yaml

import (
"bytes"
"io"

"github.com/goccy/go-yaml/internal/errors"
"golang.org/x/xerrors"
Expand Down Expand Up @@ -127,6 +128,9 @@ func Marshal(v interface{}) ([]byte, error) {
func Unmarshal(data []byte, v interface{}) error {
dec := NewDecoder(bytes.NewBuffer(data))
if err := dec.Decode(v); err != nil {
if err == io.EOF {
return nil
}
return errors.Wrapf(err, "failed to unmarshal")
}
return nil
Expand Down

0 comments on commit 68532a4

Please sign in to comment.