multi.go

Documentation: 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
View as plain text