bytes, strings: add (*Reader).WriteTo

Fixes #4031.

R=golang-dev, bradfitz, remyoudompheng, r, dave
CC=golang-dev
https://golang.org/cl/6632046

diff --git a/src/pkg/bytes/buffer_test.go b/src/pkg/bytes/buffer_test.go
index dfecea19ae39e7924bbc1ccddd19678e15934781..92e29146b32d5061d3e2ce4af69c7fb090ef3ba7 100644 (file)
--- a/src/pkg/bytes/buffer_test.go
+++ b/src/pkg/bytes/buffer_test.go
@@ -251,10 +251,10 @@ func TestReadFrom(t *testing.T) {
 func TestWriteTo(t *testing.T) {
        var buf Buffer
        for i := 3; i < 30; i += 3 {
-               s := fillBytes(t, "TestReadFrom (1)", &buf, "", 5, testBytes[0:len(testBytes)/i])
+               s := fillBytes(t, "TestWriteTo (1)", &buf, "", 5, testBytes[0:len(testBytes)/i])
                var b Buffer
                buf.WriteTo(&b)
-               empty(t, "TestReadFrom (2)", &b, s, make([]byte, len(data)))
+               empty(t, "TestWriteTo (2)", &b, s, make([]byte, len(data)))
        }
 }
 

diff --git a/src/pkg/bytes/reader.go b/src/pkg/bytes/reader.go
index a062e54ba4d0808184a6c6e074e09b437c897e9a..b34dfc11bffde916d5b963e4cff3626a2e9ddf32 100644 (file)
--- a/src/pkg/bytes/reader.go
+++ b/src/pkg/bytes/reader.go
@@ -10,7 +10,7 @@ import (
        "unicode/utf8"
 )
 
-// A Reader implements the io.Reader, io.ReaderAt, io.Seeker,
+// A Reader implements the io.Reader, io.ReaderAt, io.WriterTo, io.Seeker,
 // io.ByteScanner, and io.RuneScanner interfaces by reading from
 // a byte slice.
 // Unlike a Buffer, a Reader is read-only and supports seeking.
@@ -121,5 +121,24 @@ func (r *Reader) Seek(offset int64, whence int) (int64, error) {
        return abs, nil
 }
 
+// WriteTo implements the io.WriterTo interface.
+func (r *Reader) WriteTo(w io.Writer) (n int64, err error) {
+       r.prevRune = -1
+       if r.i >= len(r.s) {
+               return 0, io.EOF
+       }
+       b := r.s[r.i:]
+       m, err := w.Write(b)
+       if m > len(b) {
+               panic("bytes.Reader.WriteTo: invalid Write count")
+       }
+       r.i += m
+       n = int64(m)
+       if m != len(b) && err == nil {
+               err = io.ErrShortWrite
+       }
+       return
+}
+
 // NewReader returns a new Reader reading from b.
 func NewReader(b []byte) *Reader { return &Reader{b, 0, -1} }

diff --git a/src/pkg/bytes/reader_test.go b/src/pkg/bytes/reader_test.go
index 2e4b1f26e81776214c017428044d450fbe3b4df8..666881886760ea8ca6a6f78d2c9032800a14c2a4 100644 (file)
--- a/src/pkg/bytes/reader_test.go
+++ b/src/pkg/bytes/reader_test.go
@@ -86,3 +86,24 @@ func TestReaderAt(t *testing.T) {
                }
        }
 }
+
+func TestReaderWriteTo(t *testing.T) {
+       for i := 3; i < 30; i += 3 {
+               s := data[:len(data)/i]
+               r := NewReader(testBytes[:len(testBytes)/i])
+               var b Buffer
+               n, err := r.WriteTo(&b)
+               if expect := int64(len(s)); n != expect {
+                       t.Errorf("got %v; want %v", n, expect)
+               }
+               if err != nil {
+                       t.Errorf("got error = %v; want nil", err)
+               }
+               if b.String() != s {
+                       t.Errorf("got string %q; want %q", b.String(), s)
+               }
+               if r.Len() != 0 {
+                       t.Errorf("reader contains %v bytes; want 0", r.Len())
+               }
+       }
+}

diff --git a/src/pkg/strings/reader.go b/src/pkg/strings/reader.go
index 8569805552d639febfd129bfe5d62924dc57d66b..98325ce75bf659f69e14db9228e03d91dad38914 100644 (file)
--- a/src/pkg/strings/reader.go
+++ b/src/pkg/strings/reader.go
@@ -10,7 +10,7 @@ import (
        "unicode/utf8"
 )
 
-// A Reader implements the io.Reader, io.ReaderAt, io.Seeker,
+// A Reader implements the io.Reader, io.ReaderAt, io.Seeker, io.WriterTo,
 // io.ByteScanner, and io.RuneScanner interfaces by reading
 // from a string.
 type Reader struct {
@@ -120,6 +120,25 @@ func (r *Reader) Seek(offset int64, whence int) (int64, error) {
        return abs, nil
 }
 
+// WriteTo implements the io.WriterTo interface.
+func (r *Reader) WriteTo(w io.Writer) (n int64, err error) {
+       r.prevRune = -1
+       if r.i >= len(r.s) {
+               return 0, io.EOF
+       }
+       s := r.s[r.i:]
+       m, err := io.WriteString(w, s)
+       if m > len(s) {
+               panic("strings.Reader.WriteTo: invalid WriteString count")
+       }
+       r.i += m
+       n = int64(m)
+       if m != len(s) && err == nil {
+               err = io.ErrShortWrite
+       }
+       return
+}
+
 // NewReader returns a new Reader reading from s.
 // It is similar to bytes.NewBufferString but more efficient and read-only.
 func NewReader(s string) *Reader { return &Reader{s, 0, -1} }

diff --git a/src/pkg/strings/reader_test.go b/src/pkg/strings/reader_test.go
index a99ae2a0ea600fcd0c82c2837c9bccdea15148fb..bab91fc71979767871946702c2329de43eff5912 100644 (file)
--- a/src/pkg/strings/reader_test.go
+++ b/src/pkg/strings/reader_test.go
@@ -5,6 +5,7 @@
 package strings_test
 
 import (
+       "bytes"
        "fmt"
        "io"
        "os"
@@ -86,3 +87,25 @@ func TestReaderAt(t *testing.T) {
                }
        }
 }
+
+func TestWriteTo(t *testing.T) {
+       const str = "0123456789"
+       for i := 0; i < len(str); i++ {
+               s := str[i:]
+               r := strings.NewReader(s)
+               var b bytes.Buffer
+               n, err := r.WriteTo(&b)
+               if expect := int64(len(s)); n != expect {
+                       t.Errorf("got %v; want %v", n, expect)
+               }
+               if err != nil {
+                       t.Errorf("got error = %v; want nil", err)
+               }
+               if b.String() != s {
+                       t.Errorf("got string %q; want %q", b.String(), s)
+               }
+               if r.Len() != 0 {
+                       t.Errorf("reader contains %v bytes; want 0", r.Len())
+               }
+       }
+}