Preserve the original color of each token when using --with-nth with --ansi

Close #1500

4 files changed, 107 insertions, 3 deletions

diff --git a/src/ansi.go b/src/ansi.go
index d7c81d30..cc458d15 100644
--- a/src/ansi.go
+++ b/src/ansi.go
@@ -2,6 +2,7 @@ package fzf
 
 import (
 	"bytes"
+	"fmt"
 	"regexp"
 	"strconv"
 	"strings"
@@ -32,6 +33,55 @@ func (s *ansiState) equals(t *ansiState) bool {
 	return s.fg == t.fg && s.bg == t.bg && s.attr == t.attr
 }
 
+func (s *ansiState) ToString() string {
+	if !s.colored() {
+		return "\x1b[m"
+	}
+
+	ret := ""
+	if s.attr&tui.Bold > 0 {
+		ret += "1;"
+	}
+	if s.attr&tui.Dim > 0 {
+		ret += "2;"
+	}
+	if s.attr&tui.Italic > 0 {
+		ret += "3;"
+	}
+	if s.attr&tui.Underline > 0 {
+		ret += "4;"
+	}
+	if s.attr&tui.Blink > 0 {
+		ret += "5;"
+	}
+	if s.attr&tui.Reverse > 0 {
+		ret += "7;"
+	}
+	ret += toAnsiString(s.fg, 30) + toAnsiString(s.bg, 40)
+
+	return "\x1b[" + strings.TrimSuffix(ret, ";") + "m"
+}
+
+func toAnsiString(color tui.Color, offset int) string {
+	col := int(color)
+	ret := ""
+	if col == -1 {
+		ret += strconv.Itoa(offset + 9)
+	} else if col < 8 {
+		ret += strconv.Itoa(offset + col)
+	} else if col < 16 {
+		ret += strconv.Itoa(offset - 30 + 90 + col - 8)
+	} else if col < 256 {
+		ret += fmt.Sprintf("%d;5;%d", offset+8, col)
+	} else if col >= (1 << 24) {
+		r := (col >> 16) & 0xff
+		g := (col >> 8) & 0xff
+		b := col & 0xff
+		ret += fmt.Sprintf("%d;2;%d;%d;%d", offset+8, r, g, b)
+	}
+	return ret + ";"
+}
+
 var ansiRegex *regexp.Regexp
 
 func init() {
diff --git a/src/ansi_test.go b/src/ansi_test.go
index d94ae931..5acbc131 100644
--- a/src/ansi_test.go
+++ b/src/ansi_test.go
@@ -2,6 +2,7 @@ package fzf
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 
 	"github.com/junegunn/fzf/src/tui"
@@ -156,3 +157,31 @@ func TestExtractColor(t *testing.T) {
 		assert((*offsets)[1], 6, 11, 200, 100, false)
 	})
 }
+
+func TestAnsiCodeStringConversion(t *testing.T) {
+	assert := func(code string, prevState *ansiState, expected string) {
+		state := interpretCode(code, prevState)
+		if expected != state.ToString() {
+			t.Errorf("expected: %s, actual: %s",
+				strings.Replace(expected, "\x1b[", "\\x1b[", -1),
+				strings.Replace(state.ToString(), "\x1b[", "\\x1b[", -1))
+		}
+	}
+	assert("\x1b[m", nil, "\x1b[m")
+	assert("\x1b[m", &ansiState{attr: tui.Blink}, "\x1b[m")
+
+	assert("\x1b[31m", nil, "\x1b[31;49m")
+	assert("\x1b[41m", nil, "\x1b[39;41m")
+
+	assert("\x1b[92m", nil, "\x1b[92;49m")
+	assert("\x1b[102m", nil, "\x1b[39;102m")
+
+	assert("\x1b[31m", &ansiState{fg: 4, bg: 4}, "\x1b[31;44m")
+	assert("\x1b[1;2;31m", &ansiState{fg: 2, bg: -1, attr: tui.Reverse}, "\x1b[1;2;7;31;49m")
+	assert("\x1b[38;5;100;48;5;200m", nil, "\x1b[38;5;100;48;5;200m")
+	assert("\x1b[48;5;100;38;5;200m", nil, "\x1b[38;5;200;48;5;100m")
+	assert("\x1b[48;5;100;38;2;10;20;30;1m", nil, "\x1b[1;38;2;10;20;30;48;5;100m")
+	assert("\x1b[48;5;100;38;2;10;20;30;7m",
+		&ansiState{attr: tui.Dim | tui.Italic, fg: 1, bg: 1},
+		"\x1b[2;3;7;38;2;10;20;30;48;5;100m")
+}
diff --git a/src/core.go b/src/core.go
index 1653e6fd..c4b0dacb 100644
--- a/src/core.go
+++ b/src/core.go
@@ -63,12 +63,14 @@ func Run(opts *Options, revision string) {
 	ansiProcessor := func(data []byte) (util.Chars, *[]ansiOffset) {
 		return util.ToChars(data), nil
 	}
+
+	var lineAnsiState, prevLineAnsiState *ansiState
 	if opts.Ansi {
 		if opts.Theme != nil {
-			var state *ansiState
 			ansiProcessor = func(data []byte) (util.Chars, *[]ansiOffset) {
-				trimmed, offsets, newState := extractColor(string(data), state, nil)
-				state = newState
+				prevLineAnsiState = lineAnsiState
+				trimmed, offsets, newState := extractColor(string(data), lineAnsiState, nil)
+				lineAnsiState = newState
 				return util.ToChars([]byte(trimmed)), offsets
 			}
 		} else {
@@ -100,6 +102,20 @@ func Run(opts *Options, revision string) {
 	} else {
 		chunkList = NewChunkList(func(item *Item, data []byte) bool {
 			tokens := Tokenize(string(data), opts.Delimiter)
+			if opts.Ansi && len(tokens) > 1 {
+				var ansiState *ansiState
+				if prevLineAnsiState != nil {
+					ansiStateDup := *prevLineAnsiState
+					ansiState = &ansiStateDup
+				}
+				for _, token := range tokens {
+					prevAnsiState := ansiState
+					_, _, ansiState = extractColor(token.text.ToString(), ansiState, nil)
+					if prevAnsiState != nil {
+						token.text.Wrap(prevAnsiState.ToString(), "\x1b[m")
+					}
+				}
+			}
 			trans := Transform(tokens, opts.WithNth)
 			transformed := joinTokens(trans)
 			if len(header) < opts.HeaderLines {
diff --git a/src/util/chars.go b/src/util/chars.go
index ec6fca0e..35b28297 100644
--- a/src/util/chars.go
+++ b/src/util/chars.go
@@ -171,3 +171,12 @@ func (chars *Chars) CopyRunes(dest []rune) {
 	}
 	return
 }
+
+func (chars *Chars) Wrap(prefix string, suffix string) {
+	if runes := chars.optionalRunes(); runes != nil {
+		runes = append(append([]rune(prefix), runes...), []rune(suffix)...)
+		chars.slice = *(*[]byte)(unsafe.Pointer(&runes))
+	} else {
+		chars.slice = append(append([]byte(prefix), chars.slice...), []byte(suffix)...)
+	}
+}