github.com/dannin/go@v0.0.0-20161031215817-d35dfd405eaa/src/io/multi.go (about)

     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 multiReader struct {
     8  	readers []Reader
     9  }
    10  
    11  func (mr *multiReader) Read(p []byte) (n int, err error) {
    12  	for len(mr.readers) > 0 {
    13  		// Optimization to flatten nested multiReaders (Issue 13558)
    14  		if len(mr.readers) == 1 {
    15  			if r, ok := mr.readers[0].(*multiReader); ok {
    16  				mr.readers = r.readers
    17  				continue
    18  			}
    19  		}
    20  		n, err = mr.readers[0].Read(p)
    21  		if err == EOF {
    22  			mr.readers[0] = nil // permit earlier GC
    23  			mr.readers = mr.readers[1:]
    24  		}
    25  		if n > 0 || err != EOF {
    26  			if err == EOF && len(mr.readers) > 0 {
    27  				// Don't return EOF yet. More readers remain.
    28  				err = nil
    29  			}
    30  			return
    31  		}
    32  	}
    33  	return 0, EOF
    34  }
    35  
    36  // MultiReader returns a Reader that's the logical concatenation of
    37  // the provided input readers. They're read sequentially. Once all
    38  // inputs have returned EOF, Read will return EOF.  If any of the readers
    39  // return a non-nil, non-EOF error, Read will return that error.
    40  func MultiReader(readers ...Reader) Reader {
    41  	r := make([]Reader, len(readers))
    42  	copy(r, readers)
    43  	return &multiReader{r}
    44  }
    45  
    46  type multiWriter struct {
    47  	writers []Writer
    48  }
    49  
    50  func (t *multiWriter) Write(p []byte) (n int, err error) {
    51  	for _, w := range t.writers {
    52  		n, err = w.Write(p)
    53  		if err != nil {
    54  			return
    55  		}
    56  		if n != len(p) {
    57  			err = ErrShortWrite
    58  			return
    59  		}
    60  	}
    61  	return len(p), nil
    62  }
    63  
    64  var _ stringWriter = (*multiWriter)(nil)
    65  
    66  func (t *multiWriter) WriteString(s string) (n int, err error) {
    67  	var p []byte // lazily initialized if/when needed
    68  	for _, w := range t.writers {
    69  		if sw, ok := w.(stringWriter); ok {
    70  			n, err = sw.WriteString(s)
    71  		} else {
    72  			if p == nil {
    73  				p = []byte(s)
    74  			}
    75  			n, err = w.Write(p)
    76  		}
    77  		if err != nil {
    78  			return
    79  		}
    80  		if n != len(s) {
    81  			err = ErrShortWrite
    82  			return
    83  		}
    84  	}
    85  	return len(s), nil
    86  }
    87  
    88  // MultiWriter creates a writer that duplicates its writes to all the
    89  // provided writers, similar to the Unix tee(1) command.
    90  func MultiWriter(writers ...Writer) Writer {
    91  	w := make([]Writer, len(writers))
    92  	copy(w, writers)
    93  	return &multiWriter{w}
    94  }