Do not process ANSI codes in --preview output at once

Close #598

5 files changed, 19 insertions, 24 deletions

diff --git a/src/ansi.go b/src/ansi.go
index 876229f1..debba474 100644
--- a/src/ansi.go
+++ b/src/ansi.go
@@ -36,7 +36,7 @@ func init() {
 	ansiRegex = regexp.MustCompile("\x1b\\[[0-9;]*[mK]")
 }
 
-func extractColor(str string, state *ansiState) (string, []ansiOffset, *ansiState) {
+func extractColor(str string, state *ansiState, proc func(string, *ansiState) bool) (string, []ansiOffset, *ansiState) {
 	var offsets []ansiOffset
 	var output bytes.Buffer
 
@@ -46,7 +46,11 @@ func extractColor(str string, state *ansiState) (string, []ansiOffset, *ansiStat
 
 	idx := 0
 	for _, offset := range ansiRegex.FindAllStringIndex(str, -1) {
-		output.WriteString(str[idx:offset[0]])
+		prev := str[idx:offset[0]]
+		output.WriteString(prev)
+		if proc != nil && !proc(prev, state) {
+			break
+		}
 		newState := interpretCode(str[offset[0]:offset[1]], state)
 
 		if !newState.equals(state) {
@@ -77,6 +81,9 @@ func extractColor(str string, state *ansiState) (string, []ansiOffset, *ansiStat
 			(&offsets[len(offsets)-1]).offset[1] = int32(utf8.RuneCount(output.Bytes()))
 		}
 	}
+	if proc != nil {
+		proc(rest, state)
+	}
 	return output.String(), offsets, state
 }
 
diff --git a/src/ansi_test.go b/src/ansi_test.go
index e278fe98..31803f38 100644
--- a/src/ansi_test.go
+++ b/src/ansi_test.go
@@ -17,7 +17,7 @@ func TestExtractColor(t *testing.T) {
 	var state *ansiState
 	clean := "\x1b[0m"
 	check := func(assertion func(ansiOffsets []ansiOffset, state *ansiState)) {
-		output, ansiOffsets, newState := extractColor(src, state)
+		output, ansiOffsets, newState := extractColor(src, state, nil)
 		state = newState
 		if output != "hello world" {
 			t.Errorf("Invalid output: {}", output)
diff --git a/src/core.go b/src/core.go
index 24af0507..fd4bc6cd 100644
--- a/src/core.go
+++ b/src/core.go
@@ -73,7 +73,7 @@ func Run(opts *Options) {
 		if opts.Theme != nil {
 			var state *ansiState
 			ansiProcessor = func(data []byte) ([]rune, []ansiOffset) {
-				trimmed, offsets, newState := extractColor(string(data), state)
+				trimmed, offsets, newState := extractColor(string(data), state, nil)
 				state = newState
 				return []rune(trimmed), offsets
 			}
@@ -81,7 +81,7 @@ func Run(opts *Options) {
 			// When color is disabled but ansi option is given,
 			// we simply strip out ANSI codes from the input
 			ansiProcessor = func(data []byte) ([]rune, []ansiOffset) {
-				trimmed, _, _ := extractColor(string(data), nil)
+				trimmed, _, _ := extractColor(string(data), nil, nil)
 				return []rune(trimmed), nil
 			}
 		}
diff --git a/src/item.go b/src/item.go
index 66b3e4d1..06413503 100644
--- a/src/item.go
+++ b/src/item.go
@@ -128,7 +128,7 @@ func (item *Item) AsString(stripAnsi bool) string {
 func (item *Item) StringPtr(stripAnsi bool) *string {
 	if item.origText != nil {
 		if stripAnsi {
-			trimmed, _, _ := extractColor(string(*item.origText), nil)
+			trimmed, _, _ := extractColor(string(*item.origText), nil, nil)
 			return &trimmed
 		}
 		orig := string(*item.origText)
diff --git a/src/terminal.go b/src/terminal.go
index ce10ee75..18d5f0c5 100644
--- a/src/terminal.go
+++ b/src/terminal.go
@@ -556,7 +556,7 @@ func (t *Terminal) printHeader() {
 		if line >= max {
 			continue
 		}
-		trimmed, colors, newState := extractColor(lineStr, state)
+		trimmed, colors, newState := extractColor(lineStr, state, nil)
 		state = newState
 		item := &Item{
 			text:   []rune(trimmed),
@@ -730,25 +730,13 @@ func (t *Terminal) printHighlighted(item *Item, bold bool, col1 int, col2 int, c
 }
 
 func (t *Terminal) printPreview() {
-	trimmed, ansiOffsets, _ := extractColor(t.previewTxt, nil)
-	var index int32
 	t.pwindow.Erase()
-	for _, o := range ansiOffsets {
-		b := o.offset[0]
-		e := o.offset[1]
-		if b > index {
-			if !t.pwindow.Fill(trimmed[index:b]) {
-				return
-			}
-		}
-		if !t.pwindow.CFill(trimmed[b:e], o.color.fg, o.color.bg, o.color.bold) {
-			return
+	extractColor(t.previewTxt, nil, func(str string, ansi *ansiState) bool {
+		if ansi != nil && ansi.colored() {
+			return t.pwindow.CFill(str, ansi.fg, ansi.bg, ansi.bold)
 		}
-		index = e
-	}
-	if int(index) < len(trimmed) {
-		t.pwindow.Fill(trimmed[index:])
-	}
+		return t.pwindow.Fill(str)
+	})
 }
 
 func processTabs(runes []rune, prefixWidth int) (string, int) {