mirror of
https://github.com/golang/go
synced 2024-11-23 16:20:04 -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:
parent
8ceaefb74a
commit
9c514e149c
@ -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 {
|
||||
|
@ -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."
|
||||
|
Loading…
Reference in New Issue
Block a user