github.com/geraldss/go/src@v0.0.0-20210511222824-ac7d0ebfc235/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 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  }