github.com/onsi/ginkgo@v1.16.6-0.20211118180735-4e1925ba4c95/internal/writer.go (about)

     1  package internal
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"sync"
     8  )
     9  
    10  type WriterMode uint
    11  
    12  const (
    13  	WriterModeStreamAndBuffer WriterMode = iota
    14  	WriterModeBufferOnly
    15  )
    16  
    17  type WriterInterface interface {
    18  	io.Writer
    19  
    20  	Truncate()
    21  	Bytes() []byte
    22  }
    23  
    24  //Writer impplements WriterInterface and GinkgoWriterInterface
    25  type Writer struct {
    26  	buffer    *bytes.Buffer
    27  	outWriter io.Writer
    28  	lock      *sync.Mutex
    29  	mode      WriterMode
    30  
    31  	teeWriters []io.Writer
    32  }
    33  
    34  func NewWriter(outWriter io.Writer) *Writer {
    35  	return &Writer{
    36  		buffer:    &bytes.Buffer{},
    37  		lock:      &sync.Mutex{},
    38  		outWriter: outWriter,
    39  		mode:      WriterModeStreamAndBuffer,
    40  	}
    41  }
    42  
    43  func (w *Writer) SetMode(mode WriterMode) {
    44  	w.lock.Lock()
    45  	defer w.lock.Unlock()
    46  	w.mode = mode
    47  }
    48  
    49  func (w *Writer) Write(b []byte) (n int, err error) {
    50  	w.lock.Lock()
    51  	defer w.lock.Unlock()
    52  
    53  	for _, teeWriter := range w.teeWriters {
    54  		teeWriter.Write(b)
    55  	}
    56  
    57  	if w.mode == WriterModeStreamAndBuffer {
    58  		w.outWriter.Write(b)
    59  	}
    60  	return w.buffer.Write(b)
    61  }
    62  
    63  func (w *Writer) Truncate() {
    64  	w.lock.Lock()
    65  	defer w.lock.Unlock()
    66  	w.buffer.Reset()
    67  }
    68  
    69  func (w *Writer) Bytes() []byte {
    70  	w.lock.Lock()
    71  	defer w.lock.Unlock()
    72  	b := w.buffer.Bytes()
    73  	copied := make([]byte, len(b))
    74  	copy(copied, b)
    75  	return copied
    76  }
    77  
    78  //GinkgoWriterInterface
    79  func (w *Writer) TeeTo(writer io.Writer) {
    80  	w.lock.Lock()
    81  	defer w.lock.Unlock()
    82  
    83  	w.teeWriters = append(w.teeWriters, writer)
    84  }
    85  
    86  func (w *Writer) ClearTeeWriters() {
    87  	w.lock.Lock()
    88  	defer w.lock.Unlock()
    89  
    90  	w.teeWriters = []io.Writer{}
    91  }
    92  
    93  func (w *Writer) Print(a ...interface{}) {
    94  	fmt.Fprint(w, a...)
    95  }
    96  
    97  func (w *Writer) Printf(format string, a ...interface{}) {
    98  	fmt.Fprintf(w, format, a...)
    99  }
   100  
   101  func (w *Writer) Println(a ...interface{}) {
   102  	fmt.Fprintln(w, a...)
   103  }