bufio: WriteRune

also fix a printing error in the test for bytes.Buffer

R=golang-dev
CC=golang-dev
http://codereview.appspot.com/240042

2 files changed, 59 insertions, 0 deletions

diff --git a/src/pkg/bufio/bufio.go b/src/pkg/bufio/bufio.go
index 7b7f25831..9b52a363a 100644
--- a/src/pkg/bufio/bufio.go
+++ b/src/pkg/bufio/bufio.go
@@ -437,6 +437,35 @@ func (b *Writer) WriteByte(c byte) os.Error {
 	return nil
 }
 
+// WriteRune writes a single Unicode code point, returning
+// the number of bytes written and any error.
+func (b *Writer) WriteRune(rune int) (size int, err os.Error) {
+	if rune < utf8.RuneSelf {
+		err = b.WriteByte(byte(rune))
+		if err != nil {
+			return 0, err
+		}
+		return 1, nil
+	}
+	if b.err != nil {
+		return 0, b.err
+	}
+	n := b.Available()
+	if n < utf8.UTFMax {
+		if b.Flush(); b.err != nil {
+			return 0, b.err
+		}
+		n = b.Available()
+		if n < utf8.UTFMax {
+			// Can only happen if buffer is silly small.
+			return b.WriteString(string(rune))
+		}
+	}
+	size = utf8.EncodeRune(rune, b.buf[b.n:])
+	b.n += size
+	return size, nil
+}
+
 // WriteString writes a string.
 // It returns the number of bytes written.
 // If the count is less than len(s), it also returns an error explaining
diff --git a/src/pkg/bufio/bufio_test.go b/src/pkg/bufio/bufio_test.go
index 0ee8ce6b3..a19c46ec9 100644
--- a/src/pkg/bufio/bufio_test.go
+++ b/src/pkg/bufio/bufio_test.go
@@ -12,6 +12,7 @@ import (
 	"strings"
 	"testing"
 	"testing/iotest"
+	"utf8"
 )
 
 // Reads from a reader and rot13s the result.
@@ -225,6 +226,35 @@ func TestReadRune(t *testing.T) {
 	}
 }
 
+func TestReadWriteRune(t *testing.T) {
+	const NRune = 1000
+	byteBuf := new(bytes.Buffer)
+	w := NewWriter(byteBuf)
+	// Write the runes out using WriteRune
+	buf := make([]byte, utf8.UTFMax)
+	for rune := 0; rune < NRune; rune++ {
+		size := utf8.EncodeRune(rune, buf)
+		nbytes, err := w.WriteRune(rune)
+		if err != nil {
+			t.Fatalf("WriteRune(0x%x) error: %s", rune, err)
+		}
+		if nbytes != size {
+			t.Fatalf("WriteRune(0x%x) expected %d, got %d", rune, size, nbytes)
+		}
+	}
+	w.Flush()
+
+	r := NewReader(byteBuf)
+	// Read them back with ReadRune
+	for rune := 0; rune < NRune; rune++ {
+		size := utf8.EncodeRune(rune, buf)
+		nr, nbytes, err := r.ReadRune()
+		if nr != rune || nbytes != size || err != nil {
+			t.Fatalf("ReadRune(0x%x) got 0x%x,%d not 0x%x,%d (err=%s)", r, nr, nbytes, r, size, err)
+		}
+	}
+}
+
 func TestWriter(t *testing.T) {
 	var data [8192]byte