Fix: prevent prefix color change (#119)

* fix: prevent prefix color change by output

* fix: use l.currentColor

---------

Co-authored-by: ras0q <[email protected]>

run/log.go

@@ -2,25 +2,37 @@ package run
 
 import (
 	"bytes"
-	"fmt"
 	"io"
+	"regexp"
 )
 
-var newLine = byte('\n')
+var (
+	prefixColor = []byte("\033[0m")
+	delimiter   = []byte("｜ ")
+	newLine     = byte('\n')
+	colorRegexp = regexp.MustCompile(`\033\[[0-9;]*m`)
+)
 
 type prefixLogger struct {
-	w      io.Writer
-	buf    *bytes.Buffer
-	prefix []byte
+	w            io.Writer
+	buf          *bytes.Buffer
+	prefix       []byte
+	currentColor []byte
 }
 
 func newPrefixLogger(w io.Writer, prefix string) *prefixLogger {
-	streamer := &prefixLogger{
-		w:   w,
-		buf: bytes.NewBuffer([]byte("")),
-	}
+	p := make([]byte, 0, len(prefixColor)+len(prefix)+len(delimiter))
 	if prefix != "" {
-		streamer.prefix = []byte(fmt.Sprintf("%s｜ ", prefix))
+		p = append(p, prefixColor...)
+		p = append(p, []byte(prefix)...)
+		p = append(p, delimiter...)
+	}
+
+	streamer := &prefixLogger{
+		w:            w,
+		buf:          bytes.NewBuffer([]byte("")),
+		prefix:       p,
+		currentColor: []byte{},
 	}
 
 	return streamer
@@ -61,6 +73,11 @@ func (l *prefixLogger) outputLines() error {
 				if err := l.out(line); err != nil {
 					return err
 				}
+
+				colors := colorRegexp.FindAll(line, -1)
+				if len(colors) > 0 {
+					l.currentColor = colors[len(colors)-1]
+				}
 			} else {
 				// put back into buffer, it's not a complete line yet
 				//  Close() or Flush() have to be used to flush out
@@ -88,6 +105,11 @@ func (l *prefixLogger) out(p []byte) error {
 		return nil
 	}
 
-	_, err := l.w.Write(append(l.prefix, p...))
+	s := make([]byte, 0, len(l.prefix)+len(l.currentColor)+len(p))
+	s = append(s, l.prefix...)
+	s = append(s, l.currentColor...)
+	s = append(s, p...)
+
+	_, err := l.w.Write(s)
 	return err
 }

run/log_test.go

@@ -3,25 +3,9 @@ package run
 
 import (
 	"bytes"
-	"strings"
 	"testing"
 )
 
-func setExpect(prefixB []byte, s string) string {
-	prefix := string(prefixB)
-	if len(prefix) == 0 {
-		return s
-	}
-	if s == "" {
-		return s
-	}
-
-	s = strings.ReplaceAll(s, "\n", "\n"+prefix)
-	s = strings.TrimSuffix(s, prefix)
-
-	return prefix + s
-}
-
 func TestPrefixLogger_Write(t *testing.T) {
 	tests := map[string]string{"empty": "", "prefix": "prefix"}
 
@@ -35,21 +19,38 @@ func TestPrefixLogger_Write(t *testing.T) {
 			l.Write([]byte("hello world"))
 
 			// No new line so it should be empty
-			expect := setExpect(l.prefix, "")
+			expect := ""
 			if w.String() != expect {
 				t.Errorf("got %v, want %v", w.String(), expect)
 			}
 
 			// Write a new line
 			l.Write([]byte("\n"))
-			expect = setExpect(l.prefix, "hello world\n")
+			expect += string(l.prefix) + "hello world\n"
 			if w.String() != expect {
 				t.Errorf("got %v, want %v", w.String(), expect)
 			}
 
 			// Write a line with a new line
 			l.Write([]byte("foo bar\n"))
-			expect = setExpect(l.prefix, "hello world\nfoo bar\n")
+			// Each line is prefixed
+			expect += string(l.prefix) + "foo bar\n"
+			if w.String() != expect {
+				t.Errorf("got %v, want %v", w.String(), expect)
+			}
+
+			// Write a line with a new line with a red color
+			red := "\033[31m"
+			reset := "\033[m"
+			l.Write([]byte(
+				red + "this line is red\n" +
+					"this line is still red" + reset + "\n" +
+					"this line is reset\n",
+			))
+			// Each line specifies the output color
+			expect += string(l.prefix) + red + "this line is red\n"
+			expect += string(l.prefix) + red + "this line is still red" + reset + "\n"
+			expect += string(l.prefix) + reset + "this line is reset\n"
 			if w.String() != expect {
 				t.Errorf("got %v, want %v", w.String(), expect)
 			}