1 package internal
2
3 import (
4 "bytes"
5 "fmt"
6 "io"
7 "sync"
8
9 "github.com/go-logr/logr"
10 "github.com/go-logr/logr/funcr"
11 )
12
13 type WriterMode uint
14
15 const (
16 WriterModeStreamAndBuffer WriterMode = iota
17 WriterModeBufferOnly
18 )
19
20 type WriterInterface interface {
21 io.Writer
22
23 Truncate()
24 Bytes() []byte
25 Len() int
26 }
27
28
29 type Writer struct {
30 buffer *bytes.Buffer
31 outWriter io.Writer
32 lock *sync.Mutex
33 mode WriterMode
34
35 streamIndent []byte
36 indentNext bool
37
38 teeWriters []io.Writer
39 }
40
41 func NewWriter(outWriter io.Writer) *Writer {
42 return &Writer{
43 buffer: &bytes.Buffer{},
44 lock: &sync.Mutex{},
45 outWriter: outWriter,
46 mode: WriterModeStreamAndBuffer,
47 streamIndent: []byte(" "),
48 indentNext: true,
49 }
50 }
51
52 func (w *Writer) SetMode(mode WriterMode) {
53 w.lock.Lock()
54 defer w.lock.Unlock()
55 w.mode = mode
56 }
57
58 func (w *Writer) Len() int {
59 w.lock.Lock()
60 defer w.lock.Unlock()
61 return w.buffer.Len()
62 }
63
64 var newline = []byte("\n")
65
66 func (w *Writer) Write(b []byte) (n int, err error) {
67 w.lock.Lock()
68 defer w.lock.Unlock()
69
70 for _, teeWriter := range w.teeWriters {
71 teeWriter.Write(b)
72 }
73
74 if w.mode == WriterModeStreamAndBuffer {
75 line, remaining, found := []byte{}, b, false
76 for len(remaining) > 0 {
77 line, remaining, found = bytes.Cut(remaining, newline)
78 if len(line) > 0 {
79 if w.indentNext {
80 w.outWriter.Write(w.streamIndent)
81 w.indentNext = false
82 }
83 w.outWriter.Write(line)
84 }
85 if found {
86 w.outWriter.Write(newline)
87 w.indentNext = true
88 }
89 }
90 }
91 return w.buffer.Write(b)
92 }
93
94 func (w *Writer) Truncate() {
95 w.lock.Lock()
96 defer w.lock.Unlock()
97 w.buffer.Reset()
98 }
99
100 func (w *Writer) Bytes() []byte {
101 w.lock.Lock()
102 defer w.lock.Unlock()
103 b := w.buffer.Bytes()
104 copied := make([]byte, len(b))
105 copy(copied, b)
106 return copied
107 }
108
109
110 func (w *Writer) TeeTo(writer io.Writer) {
111 w.lock.Lock()
112 defer w.lock.Unlock()
113
114 w.teeWriters = append(w.teeWriters, writer)
115 }
116
117 func (w *Writer) ClearTeeWriters() {
118 w.lock.Lock()
119 defer w.lock.Unlock()
120
121 w.teeWriters = []io.Writer{}
122 }
123
124 func (w *Writer) Print(a ...interface{}) {
125 fmt.Fprint(w, a...)
126 }
127
128 func (w *Writer) Printf(format string, a ...interface{}) {
129 fmt.Fprintf(w, format, a...)
130 }
131
132 func (w *Writer) Println(a ...interface{}) {
133 fmt.Fprintln(w, a...)
134 }
135
136 func GinkgoLogrFunc(writer *Writer) logr.Logger {
137 return funcr.New(func(prefix, args string) {
138 if prefix == "" {
139 writer.Printf("%s\n", args)
140 } else {
141 writer.Printf("%s %s\n", prefix, args)
142 }
143 }, funcr.Options{})
144 }
145
View as plain text