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