Source file src/pkg/io/multi.go

     1	// Copyright 2010 The Go Authors. All rights reserved.
     2	// Use of this source code is governed by a BSD-style
     3	// license that can be found in the LICENSE file.
     4	
     5	package io
     6	
     7	type eofReader struct{}
     8	
     9	func (eofReader) Read([]byte) (int, error) {
    10		return 0, EOF
    11	}
    12	
    13	type multiReader struct {
    14		readers []Reader
    15	}
    16	
    17	func (mr *multiReader) Read(p []byte) (n int, err error) {
    18		for len(mr.readers) > 0 {
    19			// Optimization to flatten nested multiReaders (Issue 13558).
    20			if len(mr.readers) == 1 {
    21				if r, ok := mr.readers[0].(*multiReader); ok {
    22					mr.readers = r.readers
    23					continue
    24				}
    25			}
    26			n, err = mr.readers[0].Read(p)
    27			if err == EOF {
    28				// Use eofReader instead of nil to avoid nil panic
    29				// after performing flatten (Issue 18232).
    30				mr.readers[0] = eofReader{} // permit earlier GC
    31				mr.readers = mr.readers[1:]
    32			}
    33			if n > 0 || err != EOF {
    34				if err == EOF && len(mr.readers) > 0 {
    35					// Don't return EOF yet. More readers remain.
    36					err = nil
    37				}
    38				return
    39			}
    40		}
    41		return 0, EOF
    42	}
    43	
    44	// MultiReader returns a Reader that's the logical concatenation of
    45	// the provided input readers. They're read sequentially. Once all
    46	// inputs have returned EOF, Read will return EOF.  If any of the readers
    47	// return a non-nil, non-EOF error, Read will return that error.
    48	func MultiReader(readers ...Reader) Reader {
    49		r := make([]Reader, len(readers))
    50		copy(r, readers)
    51		return &multiReader{r}
    52	}
    53	
    54	type multiWriter struct {
    55		writers []Writer
    56	}
    57	
    58	func (t *multiWriter) Write(p []byte) (n int, err error) {
    59		for _, w := range t.writers {
    60			n, err = w.Write(p)
    61			if err != nil {
    62				return
    63			}
    64			if n != len(p) {
    65				err = ErrShortWrite
    66				return
    67			}
    68		}
    69		return len(p), nil
    70	}
    71	
    72	var _ StringWriter = (*multiWriter)(nil)
    73	
    74	func (t *multiWriter) WriteString(s string) (n int, err error) {
    75		var p []byte // lazily initialized if/when needed
    76		for _, w := range t.writers {
    77			if sw, ok := w.(StringWriter); ok {
    78				n, err = sw.WriteString(s)
    79			} else {
    80				if p == nil {
    81					p = []byte(s)
    82				}
    83				n, err = w.Write(p)
    84			}
    85			if err != nil {
    86				return
    87			}
    88			if n != len(s) {
    89				err = ErrShortWrite
    90				return
    91			}
    92		}
    93		return len(s), nil
    94	}
    95	
    96	// MultiWriter creates a writer that duplicates its writes to all the
    97	// provided writers, similar to the Unix tee(1) command.
    98	//
    99	// Each write is written to each listed writer, one at a time.
   100	// If a listed writer returns an error, that overall write operation
   101	// stops and returns the error; it does not continue down the list.
   102	func MultiWriter(writers ...Writer) Writer {
   103		allWriters := make([]Writer, 0, len(writers))
   104		for _, w := range writers {
   105			if mw, ok := w.(*multiWriter); ok {
   106				allWriters = append(allWriters, mw.writers...)
   107			} else {
   108				allWriters = append(allWriters, w)
   109			}
   110		}
   111		return &multiWriter{allWriters}
   112	}
   113	

View as plain text