yaml-multi-doc-iterator.go

// 
// yaml multiple document를 parsing하는데, 각 document의 타입을 미리 확정할 수 없을 때
// 각각의 document를 scanner로 분리하고 try & error로 각 document를 decoding하는 예제..
// 
package main

import (
	"bufio"
	"bytes"
	"errors"
	"testing"

	"github.com/ghodss/yaml"
	"github.com/stretchr/testify/require"
)

var data = []byte(`
---
hello2: hello2
world2: world2
---
hello1: hello1
world1: world1
`)

type helloWorld1 struct {
	Hello string `json:"hello1"`
	World string `json:"world1"`
}

type helloWorld2 struct {
	Hello string `json:"hello2"`
	World string `json:"world2"`
}

func splitYamlDoc(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}

	splitText := []byte("---\n")
	if i := bytes.Index(data, splitText); i > 0 {
		return i + len(splitText), data[0:i], nil
	}

	if atEOF {
		return len(data), data, nil
	}

	return 0, nil, nil
}

func EachYaml(data []byte, fn func([]byte) error) error {
	scanner := bufio.NewScanner(bytes.NewReader(data))
	scanner.Split(splitYamlDoc)

	for scanner.Scan() {
		if err := fn(scanner.Bytes()); err != nil {
			return err
		}
	}

	return nil
}

func TestEachYaml(t *testing.T) {
	var v1 *helloWorld1
	var v2 *helloWorld2

	err := EachYaml(data, func(doc []byte) error {
		doc = bytes.TrimSpace(doc)
		if len(doc) == 0 {
			return nil
		}

		var doc1 helloWorld1
		if err := yaml.Unmarshal(doc, &doc1); err == nil && doc1.Hello != "" {
			v1 = &doc1
			return nil
		}

		var doc2 helloWorld2
		if err := yaml.Unmarshal(doc, &doc2); err == nil && doc2.Hello != "" {
			v2 = &doc2
			return nil
		}

		return errors.New("Unknown document")
	})
	require.NoError(t, err)
	require.EqualValues(t, &helloWorld1{"hello1", "world1"}, v1)
	require.EqualValues(t, &helloWorld2{"hello2", "world2"}, v2)
}