vendor/github.com/onsi/gomega/gexec/prefixed_writer.go

   1 package gexec
   2
   3 import (
   4         "io"
   5         "sync"
   6 )
   7
   8 /*
   9 PrefixedWriter wraps an io.Writer, emiting the passed in prefix at the beginning of each new line.
  10 This can be useful when running multiple gexec.Sessions concurrently - you can prefix the log output of each
  11 session by passing in a PrefixedWriter:
  12
  13 gexec.Start(cmd, NewPrefixedWriter("[my-cmd] ", GinkgoWriter), NewPrefixedWriter("[my-cmd] ", GinkgoWriter))
  14 */
  15 type PrefixedWriter struct {
  16         prefix        []byte
  17         writer        io.Writer
  18         lock          *sync.Mutex
  19         atStartOfLine bool
  20 }
  21
  22 func NewPrefixedWriter(prefix string, writer io.Writer) *PrefixedWriter {
  23         return &PrefixedWriter{
  24                 prefix:        []byte(prefix),
  25                 writer:        writer,
  26                 lock:          &sync.Mutex{},
  27                 atStartOfLine: true,
  28         }
  29 }
  30
  31 func (w *PrefixedWriter) Write(b []byte) (int, error) {
  32         w.lock.Lock()
  33         defer w.lock.Unlock()
  34
  35         toWrite := []byte{}
  36
  37         for _, c := range b {
  38                 if w.atStartOfLine {
  39                         toWrite = append(toWrite, w.prefix...)
  40                 }
  41
  42                 toWrite = append(toWrite, c)
  43
  44                 w.atStartOfLine = c == '\n'
  45         }
  46
  47         _, err := w.writer.Write(toWrite)
  48         if err != nil {
  49                 return 0, err
  50         }
  51
  52         return len(b), nil
  53 }