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

Close #598

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
 }
 

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)

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
 			}
 		}

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)

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) {