prefix writer: simplify, preserve trailing \n
diff --git a/gexec/prefixed_writer.go b/gexec/prefixed_writer.go
index 556182b..05e695a 100644
--- a/gexec/prefixed_writer.go
+++ b/gexec/prefixed_writer.go
@@ -1,7 +1,6 @@
package gexec
import (
- "bytes"
"io"
"sync"
)
@@ -14,19 +13,18 @@
gexec.Start(cmd, NewPrefixedWriter("[my-cmd] ", GinkgoWriter), NewPrefixedWriter("[my-cmd] ", GinkgoWriter))
*/
type PrefixedWriter struct {
- prefix []byte
- writer io.Writer
- lock *sync.Mutex
- isNewLine bool
- isFirstWrite bool
+ prefix []byte
+ writer io.Writer
+ lock *sync.Mutex
+ atStartOfLine bool
}
func NewPrefixedWriter(prefix string, writer io.Writer) *PrefixedWriter {
return &PrefixedWriter{
- prefix: []byte(prefix),
- writer: writer,
- lock: &sync.Mutex{},
- isFirstWrite: true,
+ prefix: []byte(prefix),
+ writer: writer,
+ lock: &sync.Mutex{},
+ atStartOfLine: true,
}
}
@@ -34,46 +32,21 @@
w.lock.Lock()
defer w.lock.Unlock()
- newLine := []byte("\n")
- segments := bytes.Split(b, newLine)
+ toWrite := []byte{}
- if len(segments) != 0 {
- toWrite := []byte{}
- if w.isFirstWrite {
+ for _, c := range b {
+ if w.atStartOfLine {
toWrite = append(toWrite, w.prefix...)
- toWrite = append(toWrite, segments[0]...)
- w.isFirstWrite = false
- } else if w.isNewLine {
- toWrite = append(toWrite, newLine...)
- toWrite = append(toWrite, w.prefix...)
- toWrite = append(toWrite, segments[0]...)
- } else {
- toWrite = append(toWrite, segments[0]...)
}
- for i := 1; i < len(segments)-1; i++ {
- toWrite = append(toWrite, newLine...)
- toWrite = append(toWrite, w.prefix...)
- toWrite = append(toWrite, segments[i]...)
- }
+ toWrite = append(toWrite, c)
- if len(segments) > 1 {
- lastSegment := segments[len(segments)-1]
+ w.atStartOfLine = c == '\n'
+ }
- if len(lastSegment) == 0 {
- w.isNewLine = true
- } else {
- toWrite = append(toWrite, newLine...)
- toWrite = append(toWrite, w.prefix...)
- toWrite = append(toWrite, lastSegment...)
- w.isNewLine = false
- }
- }
-
- _, err := w.writer.Write(toWrite)
- if err != nil {
- return 0, err
- }
+ _, err := w.writer.Write(toWrite)
+ if err != nil {
+ return 0, err
}
return len(b), nil
diff --git a/gexec/prefixed_writer_test.go b/gexec/prefixed_writer_test.go
index 27f7487..8657d0c 100644
--- a/gexec/prefixed_writer_test.go
+++ b/gexec/prefixed_writer_test.go
@@ -2,6 +2,7 @@
import (
"bytes"
+
. "github.com/onsi/gomega/gexec"
. "github.com/onsi/ginkgo"
@@ -36,6 +37,7 @@
[p]nopqrs
[p]tuv
[p]wxyz
-[p]`))
+[p]
+`))
})
})
