implement Encoder

apic.go

@@ -80,14 +80,13 @@ func yaml_parser_set_encoding(parser *yaml_parser_t, encoding yaml_encoding_t) {
 }
 
 // Create a new emitter object.
-func yaml_emitter_initialize(emitter *yaml_emitter_t) bool {
+func yaml_emitter_initialize(emitter *yaml_emitter_t) {
 	*emitter = yaml_emitter_t{
 		buffer:     make([]byte, output_buffer_size),
 		raw_buffer: make([]byte, 0, output_raw_buffer_size),
 		states:     make([]yaml_emitter_state_t, 0, initial_stack_size),
 		events:     make([]yaml_event_t, 0, initial_queue_size),
 	}
-	return true
 }
 
 // Destroy an emitter object.
@@ -101,9 +100,10 @@ func yaml_string_write_handler(emitter *yaml_emitter_t, buffer []byte) error {
 	return nil
 }
 
-// File write handler.
-func yaml_file_write_handler(emitter *yaml_emitter_t, buffer []byte) error {
-	_, err := emitter.output_file.Write(buffer)
+// yaml_writer_write_handler uses emitter.output_writer to write the
+// emitted text.
+func yaml_writer_write_handler(emitter *yaml_emitter_t, buffer []byte) error {
+	_, err := emitter.output_writer.Write(buffer)
 	return err
 }
 
@@ -117,12 +117,12 @@ func yaml_emitter_set_output_string(emitter *yaml_emitter_t, output_buffer *[]by
 }
 
 // Set a file output.
-func yaml_emitter_set_output_file(emitter *yaml_emitter_t, file io.Writer) {
+func yaml_emitter_set_output_writer(emitter *yaml_emitter_t, w io.Writer) {
 	if emitter.write_handler != nil {
 		panic("must set the output target only once")
 	}
-	emitter.write_handler = yaml_file_write_handler
-	emitter.output_file = file
+	emitter.write_handler = yaml_writer_write_handler
+	emitter.output_writer = w
 }
 
 // Set the output encoding.
@@ -251,41 +251,41 @@ func yaml_emitter_set_break(emitter *yaml_emitter_t, line_break yaml_break_t) {
 //
 
 // Create STREAM-START.
-func yaml_stream_start_event_initialize(event *yaml_event_t, encoding yaml_encoding_t) bool {
+func yaml_stream_start_event_initialize(event *yaml_event_t, encoding yaml_encoding_t) {
 	*event = yaml_event_t{
 		typ:      yaml_STREAM_START_EVENT,
 		encoding: encoding,
 	}
-	return true
 }
 
 // Create STREAM-END.
-func yaml_stream_end_event_initialize(event *yaml_event_t) bool {
+func yaml_stream_end_event_initialize(event *yaml_event_t) {
 	*event = yaml_event_t{
 		typ: yaml_STREAM_END_EVENT,
 	}
-	return true
 }
 
 // Create DOCUMENT-START.
-func yaml_document_start_event_initialize(event *yaml_event_t, version_directive *yaml_version_directive_t,
-	tag_directives []yaml_tag_directive_t, implicit bool) bool {
+func yaml_document_start_event_initialize(
+	event *yaml_event_t,
+	version_directive *yaml_version_directive_t,
+	tag_directives []yaml_tag_directive_t,
+	implicit bool,
+) {
 	*event = yaml_event_t{
 		typ:               yaml_DOCUMENT_START_EVENT,
 		version_directive: version_directive,
 		tag_directives:    tag_directives,
 		implicit:          implicit,
 	}
-	return true
 }
 
 // Create DOCUMENT-END.
-func yaml_document_end_event_initialize(event *yaml_event_t, implicit bool) bool {
+func yaml_document_end_event_initialize(event *yaml_event_t, implicit bool) {
 	*event = yaml_event_t{
 		typ:      yaml_DOCUMENT_END_EVENT,
 		implicit: implicit,
 	}
-	return true
 }
 
 ///*
@@ -347,7 +347,7 @@ func yaml_sequence_end_event_initialize(event *yaml_event_t) bool {
 }
 
 // Create MAPPING-START.
-func yaml_mapping_start_event_initialize(event *yaml_event_t, anchor, tag []byte, implicit bool, style yaml_mapping_style_t) bool {
+func yaml_mapping_start_event_initialize(event *yaml_event_t, anchor, tag []byte, implicit bool, style yaml_mapping_style_t) {
 	*event = yaml_event_t{
 		typ:      yaml_MAPPING_START_EVENT,
 		anchor:   anchor,
@@ -355,15 +355,13 @@ func yaml_mapping_start_event_initialize(event *yaml_event_t, anchor, tag []byte
 		implicit: implicit,
 		style:    yaml_style_t(style),
 	}
-	return true
 }
 
 // Create MAPPING-END.
-func yaml_mapping_end_event_initialize(event *yaml_event_t) bool {
+func yaml_mapping_end_event_initialize(event *yaml_event_t) {
 	*event = yaml_event_t{
 		typ: yaml_MAPPING_END_EVENT,
 	}
-	return true
 }
 
 // Destroy an event object.

emitterc.go

@@ -2,6 +2,7 @@ package yaml
 
 import (
 	"bytes"
+	"fmt"
 )
 
 // Flush the buffer if needed.
@@ -664,7 +665,7 @@ func yaml_emitter_emit_node(emitter *yaml_emitter_t, event *yaml_event_t,
 		return yaml_emitter_emit_mapping_start(emitter, event)
 	default:
 		return yaml_emitter_set_emitter_error(emitter,
-			"expected SCALAR, SEQUENCE-START, MAPPING-START, or ALIAS")
+			fmt.Sprintf("expected SCALAR, SEQUENCE-START, MAPPING-START, or ALIAS, but got %v", event.typ))
 	}
 }
 

encode.go

@@ -3,6 +3,7 @@ package yaml
 import (
 	"encoding"
 	"fmt"
+	"io"
 	"reflect"
 	"regexp"
 	"sort"
@@ -16,25 +17,39 @@ type encoder struct {
 	event   yaml_event_t
 	out     []byte
 	flow    bool
+	// doneInit holds whether the initial stream_start_event has been
+	// emitted.
+	doneInit bool
 }
 
-func newEncoder() (e *encoder) {
-	e = &encoder{}
-	e.must(yaml_emitter_initialize(&e.emitter))
+func newEncoder() *encoder {
+	e := &encoder{}
+	yaml_emitter_initialize(&e.emitter)
 	yaml_emitter_set_output_string(&e.emitter, &e.out)
 	yaml_emitter_set_unicode(&e.emitter, true)
-	e.must(yaml_stream_start_event_initialize(&e.event, yaml_UTF8_ENCODING))
-	e.emit()
-	e.must(yaml_document_start_event_initialize(&e.event, nil, nil, true))
-	e.emit()
 	return e
 }
 
-func (e *encoder) finish() {
-	e.must(yaml_document_end_event_initialize(&e.event, true))
+func newEncoderWithWriter(w io.Writer) *encoder {
+	e := &encoder{}
+	yaml_emitter_initialize(&e.emitter)
+	yaml_emitter_set_output_writer(&e.emitter, w)
+	yaml_emitter_set_unicode(&e.emitter, true)
+	return e
+}
+
+func (e *encoder) init() {
+	if e.doneInit {
+		return
+	}
+	yaml_stream_start_event_initialize(&e.event, yaml_UTF8_ENCODING)
 	e.emit()
+	e.doneInit = true
+}
+
+func (e *encoder) finish() {
 	e.emitter.open_ended = false
-	e.must(yaml_stream_end_event_initialize(&e.event))
+	yaml_stream_end_event_initialize(&e.event)
 	e.emit()
 }
 
@@ -44,9 +59,7 @@ func (e *encoder) destroy() {
 
 func (e *encoder) emit() {
 	// This will internally delete the e.event value.
-	if !yaml_emitter_emit(&e.emitter, &e.event) && e.event.typ != yaml_DOCUMENT_END_EVENT && e.event.typ != yaml_STREAM_END_EVENT {
-		e.must(false)
-	}
+	e.must(yaml_emitter_emit(&e.emitter, &e.event))
 }
 
 func (e *encoder) must(ok bool) {
@@ -59,6 +72,15 @@ func (e *encoder) must(ok bool) {
 	}
 }
 
+func (e *encoder) marshalDoc(tag string, in reflect.Value) {
+	e.init()
+	yaml_document_start_event_initialize(&e.event, nil, nil, true)
+	e.emit()
+	e.marshal(tag, in)
+	yaml_document_end_event_initialize(&e.event, true)
+	e.emit()
+}
+
 func (e *encoder) marshal(tag string, in reflect.Value) {
 	if !in.IsValid() {
 		e.nilv()
@@ -191,10 +213,10 @@ func (e *encoder) mappingv(tag string, f func()) {
 		e.flow = false
 		style = yaml_FLOW_MAPPING_STYLE
 	}
-	e.must(yaml_mapping_start_event_initialize(&e.event, nil, []byte(tag), implicit, style))
+	yaml_mapping_start_event_initialize(&e.event, nil, []byte(tag), implicit, style)
 	e.emit()
 	f()
-	e.must(yaml_mapping_end_event_initialize(&e.event))
+	yaml_mapping_end_event_initialize(&e.event)
 	e.emit()
 }
 

encode_test.go

@@ -1,16 +1,18 @@
 package yaml_test
 
 import (
+	"bytes"
 	"fmt"
 	"math"
 	"strconv"
 	"strings"
 	"time"
 
-	. "gopkg.in/check.v1"
-	"gopkg.in/yaml.v2"
 	"net"
 	"os"
+
+	. "gopkg.in/check.v1"
+	"gopkg.in/yaml.v2"
 )
 
 var marshalIntTest = 123
@@ -327,13 +329,51 @@ var marshalTests = []struct {
 func (s *S) TestMarshal(c *C) {
 	defer os.Setenv("TZ", os.Getenv("TZ"))
 	os.Setenv("TZ", "UTC")
-	for _, item := range marshalTests {
+	for i, item := range marshalTests {
+		c.Logf("test %d. %q", i, item.data)
 		data, err := yaml.Marshal(item.value)
 		c.Assert(err, IsNil)
 		c.Assert(string(data), Equals, item.data)
 	}
 }
 
+func (s *S) TestEncoderSingleDocument(c *C) {
+	for i, item := range marshalTests {
+		c.Logf("test %d. %q", i, item.data)
+		var buf bytes.Buffer
+		enc := yaml.NewEncoder(&buf)
+		err := enc.Encode(item.value)
+		c.Assert(err, Equals, nil)
+		err = enc.Close()
+		c.Assert(err, Equals, nil)
+		c.Assert(buf.String(), Equals, item.data)
+	}
+}
+
+func (s *S) TestEncoderMultipleDocuments(c *C) {
+	var buf bytes.Buffer
+	enc := yaml.NewEncoder(&buf)
+	err := enc.Encode(map[string]string{"a": "b"})
+	c.Assert(err, Equals, nil)
+	err = enc.Encode(map[string]string{"c": "d"})
+	c.Assert(err, Equals, nil)
+	err = enc.Close()
+	c.Assert(err, Equals, nil)
+	c.Assert(buf.String(), Equals, "a: b\n---\nc: d\n")
+}
+
+func (s *S) TestEncoderWriteError(c *C) {
+	enc := yaml.NewEncoder(errorWriter{})
+	err := enc.Encode(map[string]string{"a": "b"})
+	c.Assert(err, ErrorMatches, `yaml: write error: some write error`) // Data not flushed yet
+}
+
+type errorWriter struct{}
+
+func (errorWriter) Write([]byte) (int, error) {
+	return 0, fmt.Errorf("some write error")
+}
+
 var marshalErrorTests = []struct {
 	value interface{}
 	error string

writerc.go

@@ -18,72 +18,9 @@ func yaml_emitter_flush(emitter *yaml_emitter_t) bool {
 		return true
 	}
 
-	// If the output encoding is UTF-8, we don't need to recode the buffer.
-	if emitter.encoding == yaml_UTF8_ENCODING {
-		if err := emitter.write_handler(emitter, emitter.buffer[:emitter.buffer_pos]); err != nil {
-			return yaml_emitter_set_writer_error(emitter, "write error: "+err.Error())
-		}
-		emitter.buffer_pos = 0
-		return true
-	}
-
-	// Recode the buffer into the raw buffer.
-	var low, high int
-	if emitter.encoding == yaml_UTF16LE_ENCODING {
-		low, high = 0, 1
-	} else {
-		high, low = 1, 0
-	}
-
-	pos := 0
-	for pos < emitter.buffer_pos {
-		// See the "reader.c" code for more details on UTF-8 encoding.  Note
-		// that we assume that the buffer contains a valid UTF-8 sequence.
-
-		// Read the next UTF-8 character.
-		octet := emitter.buffer[pos]
-
-		var w int
-		var value rune
-		switch {
-		case octet&0x80 == 0x00:
-			w, value = 1, rune(octet&0x7F)
-		case octet&0xE0 == 0xC0:
-			w, value = 2, rune(octet&0x1F)
-		case octet&0xF0 == 0xE0:
-			w, value = 3, rune(octet&0x0F)
-		case octet&0xF8 == 0xF0:
-			w, value = 4, rune(octet&0x07)
-		}
-		for k := 1; k < w; k++ {
-			octet = emitter.buffer[pos+k]
-			value = (value << 6) + (rune(octet) & 0x3F)
-		}
-		pos += w
-
-		// Write the character.
-		if value < 0x10000 {
-			var b [2]byte
-			b[high] = byte(value >> 8)
-			b[low] = byte(value & 0xFF)
-			emitter.raw_buffer = append(emitter.raw_buffer, b[0], b[1])
-		} else {
-			// Write the character using a surrogate pair (check "reader.c").
-			var b [4]byte
-			value -= 0x10000
-			b[high] = byte(0xD8 + (value >> 18))
-			b[low] = byte((value >> 10) & 0xFF)
-			b[high+2] = byte(0xDC + ((value >> 8) & 0xFF))
-			b[low+2] = byte(value & 0xFF)
-			emitter.raw_buffer = append(emitter.raw_buffer, b[0], b[1], b[2], b[3])
-		}
-	}
-
-	// Write the raw buffer.
-	if err := emitter.write_handler(emitter, emitter.raw_buffer); err != nil {
+	if err := emitter.write_handler(emitter, emitter.buffer[:emitter.buffer_pos]); err != nil {
 		return yaml_emitter_set_writer_error(emitter, "write error: "+err.Error())
 	}
 	emitter.buffer_pos = 0
-	emitter.raw_buffer = emitter.raw_buffer[:0]
 	return true
 }

yaml.go

@@ -196,12 +196,47 @@ func Marshal(in interface{}) (out []byte, err error) {
 	defer handleErr(&err)
 	e := newEncoder()
 	defer e.destroy()
-	e.marshal("", reflect.ValueOf(in))
+	e.marshalDoc("", reflect.ValueOf(in))
 	e.finish()
 	out = e.out
 	return
 }
 
+// An Encoder writes YAML values to an output stream.
+type Encoder struct {
+	encoder *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: newEncoderWithWriter(w),
+	}
+}
+
+// Encode writes the YAML encoding of v 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(v interface{}) (err error) {
+	defer handleErr(&err)
+	e.encoder.marshalDoc("", reflect.ValueOf(v))
+	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) {
+	defer handleErr(&err)
+	e.encoder.finish()
+	return nil
+}
+
 func handleErr(err *error) {
 	if v := recover(); v != nil {
 		if e, ok := v.(yamlError); ok {

yamlh.go

@@ -632,7 +632,7 @@ type yaml_emitter_t struct {
 	write_handler yaml_write_handler_t // Write handler.
 
 	output_buffer *[]byte   // String output data.
-	output_file   io.Writer // File output data.
+	output_writer io.Writer // File output data.
 
 	buffer     []byte // The working buffer.
 	buffer_pos int    // The current position of the buffer.