Home | History | Annotate | Download | only in io 
      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