1
0
mirror of https://github.com/golang/go synced 2024-11-12 09:10:21 -07:00

io: add WriteString support to MultiWriter

Fixes #11805

Change-Id: I081e16b869dc706bd847ee645bb902bc671c123f
Reviewed-on: https://go-review.googlesource.com/12485
Reviewed-by: Josh Bleecher Snyder <josharian@gmail.com>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
This commit is contained in:
Brad Fitzpatrick 2015-07-21 12:24:33 -07:00
parent 8ceaefb74a
commit 9c514e149c
2 changed files with 76 additions and 1 deletions

View File

@ -52,6 +52,30 @@ func (t *multiWriter) Write(p []byte) (n int, err error) {
return len(p), nil
}
var _ stringWriter = (*multiWriter)(nil)
func (t *multiWriter) WriteString(s string) (n int, err error) {
var p []byte // lazily initialized if/when needed
for _, w := range t.writers {
if sw, ok := w.(stringWriter); ok {
n, err = sw.WriteString(s)
} else {
if p == nil {
p = []byte(s)
}
n, err = w.Write(p)
}
if err != nil {
return
}
if n != len(s) {
err = ErrShortWrite
return
}
}
return len(s), nil
}
// MultiWriter creates a writer that duplicates its writes to all the
// provided writers, similar to the Unix tee(1) command.
func MultiWriter(writers ...Writer) Writer {

View File

@ -62,8 +62,59 @@ func TestMultiReader(t *testing.T) {
}
func TestMultiWriter(t *testing.T) {
sha1 := sha1.New()
sink := new(bytes.Buffer)
// Hide bytes.Buffer's WriteString method:
testMultiWriter(t, struct {
Writer
fmt.Stringer
}{sink, sink})
}
func TestMultiWriter_String(t *testing.T) {
testMultiWriter(t, new(bytes.Buffer))
}
// test that a multiWriter.WriteString calls results in at most 1 allocation,
// even if multiple targets don't support WriteString.
func TestMultiWriter_WriteStringSingleAlloc(t *testing.T) {
var sink1, sink2 bytes.Buffer
type simpleWriter struct { // hide bytes.Buffer's WriteString
Writer
}
mw := MultiWriter(simpleWriter{&sink1}, simpleWriter{&sink2})
allocs := int(testing.AllocsPerRun(1000, func() {
WriteString(mw, "foo")
}))
if allocs != 1 {
t.Errorf("num allocations = %d; want 1", allocs)
}
}
type writeStringChecker struct{ called bool }
func (c *writeStringChecker) WriteString(s string) (n int, err error) {
c.called = true
return len(s), nil
}
func (c *writeStringChecker) Write(p []byte) (n int, err error) {
return len(p), nil
}
func TestMultiWriter_StringCheckCall(t *testing.T) {
var c writeStringChecker
mw := MultiWriter(&c)
WriteString(mw, "foo")
if !c.called {
t.Error("did not see WriteString call to writeStringChecker")
}
}
func testMultiWriter(t *testing.T, sink interface {
Writer
fmt.Stringer
}) {
sha1 := sha1.New()
mw := MultiWriter(sha1, sink)
sourceString := "My input text."