libgo/go/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 }