2022-08-24 02:24:25 +00:00
|
|
|
package internal
|
|
|
|
|
|
|
|
import (
|
|
|
|
"bytes"
|
|
|
|
"fmt"
|
|
|
|
"io"
|
|
|
|
"sync"
|
2022-10-24 20:14:32 +00:00
|
|
|
|
|
|
|
"github.com/go-logr/logr"
|
|
|
|
"github.com/go-logr/logr/funcr"
|
2022-08-24 02:24:25 +00:00
|
|
|
)
|
|
|
|
|
|
|
|
type WriterMode uint
|
|
|
|
|
|
|
|
const (
|
|
|
|
WriterModeStreamAndBuffer WriterMode = iota
|
|
|
|
WriterModeBufferOnly
|
|
|
|
)
|
|
|
|
|
|
|
|
type WriterInterface interface {
|
|
|
|
io.Writer
|
|
|
|
|
|
|
|
Truncate()
|
|
|
|
Bytes() []byte
|
2023-06-02 15:01:18 +00:00
|
|
|
Len() int
|
2022-08-24 02:24:25 +00:00
|
|
|
}
|
|
|
|
|
2023-06-02 15:01:18 +00:00
|
|
|
// Writer implements WriterInterface and GinkgoWriterInterface
|
2022-08-24 02:24:25 +00:00
|
|
|
type Writer struct {
|
|
|
|
buffer *bytes.Buffer
|
|
|
|
outWriter io.Writer
|
|
|
|
lock *sync.Mutex
|
|
|
|
mode WriterMode
|
|
|
|
|
2023-06-02 15:01:18 +00:00
|
|
|
streamIndent []byte
|
|
|
|
indentNext bool
|
|
|
|
|
2022-08-24 02:24:25 +00:00
|
|
|
teeWriters []io.Writer
|
|
|
|
}
|
|
|
|
|
|
|
|
func NewWriter(outWriter io.Writer) *Writer {
|
|
|
|
return &Writer{
|
2023-06-02 15:01:18 +00:00
|
|
|
buffer: &bytes.Buffer{},
|
|
|
|
lock: &sync.Mutex{},
|
|
|
|
outWriter: outWriter,
|
|
|
|
mode: WriterModeStreamAndBuffer,
|
|
|
|
streamIndent: []byte(" "),
|
|
|
|
indentNext: true,
|
2022-08-24 02:24:25 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) SetMode(mode WriterMode) {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
w.mode = mode
|
|
|
|
}
|
|
|
|
|
2023-06-02 15:01:18 +00:00
|
|
|
func (w *Writer) Len() int {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
return w.buffer.Len()
|
|
|
|
}
|
|
|
|
|
|
|
|
var newline = []byte("\n")
|
|
|
|
|
2022-08-24 02:24:25 +00:00
|
|
|
func (w *Writer) Write(b []byte) (n int, err error) {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
|
|
|
|
for _, teeWriter := range w.teeWriters {
|
|
|
|
teeWriter.Write(b)
|
|
|
|
}
|
|
|
|
|
|
|
|
if w.mode == WriterModeStreamAndBuffer {
|
2023-06-02 15:01:18 +00:00
|
|
|
line, remaining, found := []byte{}, b, false
|
|
|
|
for len(remaining) > 0 {
|
|
|
|
line, remaining, found = bytes.Cut(remaining, newline)
|
|
|
|
if len(line) > 0 {
|
|
|
|
if w.indentNext {
|
|
|
|
w.outWriter.Write(w.streamIndent)
|
|
|
|
w.indentNext = false
|
|
|
|
}
|
|
|
|
w.outWriter.Write(line)
|
|
|
|
}
|
|
|
|
if found {
|
|
|
|
w.outWriter.Write(newline)
|
|
|
|
w.indentNext = true
|
|
|
|
}
|
|
|
|
}
|
2022-08-24 02:24:25 +00:00
|
|
|
}
|
|
|
|
return w.buffer.Write(b)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) Truncate() {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
w.buffer.Reset()
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) Bytes() []byte {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
b := w.buffer.Bytes()
|
|
|
|
copied := make([]byte, len(b))
|
|
|
|
copy(copied, b)
|
|
|
|
return copied
|
|
|
|
}
|
|
|
|
|
2023-06-02 15:01:18 +00:00
|
|
|
// GinkgoWriterInterface
|
2022-08-24 02:24:25 +00:00
|
|
|
func (w *Writer) TeeTo(writer io.Writer) {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
|
|
|
|
w.teeWriters = append(w.teeWriters, writer)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) ClearTeeWriters() {
|
|
|
|
w.lock.Lock()
|
|
|
|
defer w.lock.Unlock()
|
|
|
|
|
|
|
|
w.teeWriters = []io.Writer{}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) Print(a ...interface{}) {
|
|
|
|
fmt.Fprint(w, a...)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) Printf(format string, a ...interface{}) {
|
|
|
|
fmt.Fprintf(w, format, a...)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) Println(a ...interface{}) {
|
|
|
|
fmt.Fprintln(w, a...)
|
|
|
|
}
|
2022-10-24 20:14:32 +00:00
|
|
|
|
|
|
|
func GinkgoLogrFunc(writer *Writer) logr.Logger {
|
|
|
|
return funcr.New(func(prefix, args string) {
|
2023-09-25 20:31:44 +00:00
|
|
|
if prefix == "" {
|
|
|
|
writer.Printf("%s\n", args)
|
|
|
|
} else {
|
|
|
|
writer.Printf("%s %s\n", prefix, args)
|
|
|
|
}
|
2022-10-24 20:14:32 +00:00
|
|
|
}, funcr.Options{})
|
|
|
|
}
|