Make --accept-nth and --with-nth support templates

5 files changed, 82 insertions, 43 deletions

diff --git a/src/core.go b/src/core.go
index 08d9e868..939910b3 100644
--- a/src/core.go
+++ b/src/core.go
@@ -96,7 +96,7 @@ func Run(opts *Options) (int, error) {
 	var chunkList *ChunkList
 	var itemIndex int32
 	header := make([]string, 0, opts.HeaderLines)
-	if len(opts.WithNth) == 0 {
+	if opts.WithNth == nil {
 		chunkList = NewChunkList(cache, func(item *Item, data []byte) bool {
 			if len(header) < opts.HeaderLines {
 				header = append(header, byteString(data))
@@ -109,6 +109,7 @@ func Run(opts *Options) (int, error) {
 			return true
 		})
 	} else {
+		nthTransformer := opts.WithNth(opts.Delimiter)
 		chunkList = NewChunkList(cache, func(item *Item, data []byte) bool {
 			tokens := Tokenize(byteString(data), opts.Delimiter)
 			if opts.Ansi && opts.Theme.Colored && len(tokens) > 1 {
@@ -127,15 +128,13 @@ func Run(opts *Options) (int, error) {
 					}
 				}
 			}
-			trans := Transform(tokens, opts.WithNth)
-			transformed := JoinTokens(trans)
+			transformed := nthTransformer(tokens)
 			if len(header) < opts.HeaderLines {
 				header = append(header, transformed)
 				eventBox.Set(EvtHeader, header)
 				return false
 			}
 			item.text, item.colors = ansiProcessor(stringBytes(transformed))
-			item.text.TrimTrailingWhitespaces()
 			item.text.Index = itemIndex
 			item.origText = &data
 			itemIndex++
diff --git a/src/options.go b/src/options.go
index cad2936e..80079a4c 100644
--- a/src/options.go
+++ b/src/options.go
@@ -544,8 +544,8 @@ type Options struct {
 	Case              Case
 	Normalize         bool
 	Nth               []Range
-	WithNth           []Range
-	AcceptNth         []Range
+	WithNth           func(Delimiter) func([]Token) string
+	AcceptNth         func(Delimiter) func([]Token) string
 	Delimiter         Delimiter
 	Sort              int
 	Track             trackOption
@@ -667,8 +667,6 @@ func defaultOptions() *Options {
 		Case:         CaseSmart,
 		Normalize:    true,
 		Nth:          make([]Range, 0),
-		WithNth:      make([]Range, 0),
-		AcceptNth:    make([]Range, 0),
 		Delimiter:    Delimiter{},
 		Sort:         1000,
 		Track:        trackDisabled,
@@ -771,6 +769,62 @@ func splitNth(str string) ([]Range, error) {
 	return ranges, nil
 }
 
+func nthTransformer(str string) (func(Delimiter) func([]Token) string, error) {
+	// ^[0-9,-.]+$"
+	if match, _ := regexp.MatchString("^[0-9,-.]+$", str); match {
+		nth, err := splitNth(str)
+		if err != nil {
+			return nil, err
+		}
+		return func(Delimiter) func([]Token) string {
+			return func(tokens []Token) string {
+				return JoinTokens(Transform(tokens, nth))
+			}
+		}, nil
+	}
+
+	// {...} {...} ...
+	placeholder := regexp.MustCompile("{[0-9,-.]+}")
+	indexes := placeholder.FindAllStringIndex(str, -1)
+	if indexes == nil {
+		return nil, errors.New("template should include at least 1 placeholder: " + str)
+	}
+
+	type NthParts struct {
+		str string
+		nth []Range
+	}
+
+	parts := make([]NthParts, len(indexes))
+	idx := 0
+	for _, index := range indexes {
+		if idx < index[0] {
+			parts = append(parts, NthParts{str: str[idx:index[0]]})
+		}
+		if nth, err := splitNth(str[index[0]+1 : index[1]-1]); err == nil {
+			parts = append(parts, NthParts{nth: nth})
+		}
+		idx = index[1]
+	}
+	if idx < len(str) {
+		parts = append(parts, NthParts{str: str[idx:]})
+	}
+
+	return func(delimiter Delimiter) func([]Token) string {
+		return func(tokens []Token) string {
+			str := ""
+			for _, holder := range parts {
+				if holder.nth != nil {
+					str += StripLastDelimiter(JoinTokens(Transform(tokens, holder.nth)), delimiter)
+				} else {
+					str += holder.str
+				}
+			}
+			return str
+		}
+	}, nil
+}
+
 func delimiterRegexp(str string) Delimiter {
 	// Special handling of \t
 	str = strings.ReplaceAll(str, "\\t", "\t")
@@ -2387,7 +2441,7 @@ func parseOptions(index *int, opts *Options, allArgs []string) error {
 			if err != nil {
 				return err
 			}
-			if opts.WithNth, err = splitNth(str); err != nil {
+			if opts.WithNth, err = nthTransformer(str); err != nil {
 				return err
 			}
 		case "--accept-nth":
@@ -2395,7 +2449,7 @@ func parseOptions(index *int, opts *Options, allArgs []string) error {
 			if err != nil {
 				return err
 			}
-			if opts.AcceptNth, err = splitNth(str); err != nil {
+			if opts.AcceptNth, err = nthTransformer(str); err != nil {
 				return err
 			}
 		case "-s", "--sort":
diff --git a/src/terminal.go b/src/terminal.go
index 273f2650..9a4abf86 100644
--- a/src/terminal.go
+++ b/src/terminal.go
@@ -305,7 +305,7 @@ type Terminal struct {
 	nthAttr            tui.Attr
 	nth                []Range
 	nthCurrent         []Range
-	acceptNth          []Range
+	acceptNth          func([]Token) string
 	tabstop            int
 	margin             [4]sizeSpec
 	padding            [4]sizeSpec
@@ -919,7 +919,6 @@ func NewTerminal(opts *Options, eventBox *util.EventBox, executor *util.Executor
 		nthAttr:            opts.Theme.Nth.Attr,
 		nth:                opts.Nth,
 		nthCurrent:         opts.Nth,
-		acceptNth:          opts.AcceptNth,
 		tabstop:            opts.Tabstop,
 		hasStartActions:    false,
 		hasResultActions:   false,
@@ -961,6 +960,9 @@ func NewTerminal(opts *Options, eventBox *util.EventBox, executor *util.Executor
 		lastAction:         actStart,
 		lastFocus:          minItem.Index(),
 		numLinesCache:      make(map[int32]numLinesCacheValue)}
+	if opts.AcceptNth != nil {
+		t.acceptNth = opts.AcceptNth(t.delimiter)
+	}
 
 	// This should be called before accessing tui.Color*
 	tui.InitTheme(opts.Theme, renderer.DefaultTheme(), opts.Black, opts.InputBorderShape.Visible(), opts.HeaderBorderShape.Visible())
@@ -1570,9 +1572,11 @@ func (t *Terminal) output() bool {
 	transform := func(item *Item) string {
 		return item.AsString(t.ansi)
 	}
-	if len(t.acceptNth) > 0 {
+	if t.acceptNth != nil {
 		transform = func(item *Item) string {
-			return JoinTokens(StripLastDelimiter(Transform(Tokenize(item.AsString(t.ansi), t.delimiter), t.acceptNth), t.delimiter))
+			tokens := Tokenize(item.AsString(t.ansi), t.delimiter)
+			transformed := t.acceptNth(tokens)
+			return StripLastDelimiter(transformed, t.delimiter)
 		}
 	}
 	found := len(t.selected) > 0
diff --git a/src/tokenizer.go b/src/tokenizer.go
index 057d7405..aaddd17d 100644
--- a/src/tokenizer.go
+++ b/src/tokenizer.go
@@ -6,6 +6,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"unicode"
 
 	"github.com/junegunn/fzf/src/util"
 )
@@ -211,32 +212,18 @@ func Tokenize(text string, delimiter Delimiter) []Token {
 	return withPrefixLengths(tokens, 0)
 }
 
-// StripLastDelimiter removes the trailing delimiter and whitespaces from the
-// last token.
-func StripLastDelimiter(tokens []Token, delimiter Delimiter) []Token {
-	if len(tokens) == 0 {
-		return tokens
-	}
-
-	lastToken := tokens[len(tokens)-1]
-
-	if delimiter.str == nil && delimiter.regex == nil {
-		lastToken.text.TrimTrailingWhitespaces()
-	} else {
-		if delimiter.str != nil {
-			lastToken.text.TrimSuffix([]rune(*delimiter.str))
-		} else if delimiter.regex != nil {
-			str := lastToken.text.ToString()
-			locs := delimiter.regex.FindAllStringIndex(str, -1)
-			if len(locs) > 0 {
-				lastLoc := locs[len(locs)-1]
-				lastToken.text.SliceRight(lastLoc[0])
-			}
+// StripLastDelimiter removes the trailing delimiter and whitespaces
+func StripLastDelimiter(str string, delimiter Delimiter) string {
+	if delimiter.str != nil {
+		str = strings.TrimSuffix(str, *delimiter.str)
+	} else if delimiter.regex != nil {
+		locs := delimiter.regex.FindAllStringIndex(str, -1)
+		if len(locs) > 0 {
+			lastLoc := locs[len(locs)-1]
+			str = str[:lastLoc[0]]
 		}
-		lastToken.text.TrimTrailingWhitespaces()
 	}
-
-	return tokens
+	return strings.TrimRightFunc(str, unicode.IsSpace)
 }
 
 // JoinTokens concatenates the tokens into a single string
diff --git a/src/util/chars.go b/src/util/chars.go
index dd037caa..adde02a6 100644
--- a/src/util/chars.go
+++ b/src/util/chars.go
@@ -184,11 +184,6 @@ func (chars *Chars) TrailingWhitespaces() int {
 	return whitespaces
 }
 
-func (chars *Chars) TrimTrailingWhitespaces() {
-	whitespaces := chars.TrailingWhitespaces()
-	chars.slice = chars.slice[0 : len(chars.slice)-whitespaces]
-}
-
 func (chars *Chars) TrimSuffix(runes []rune) {
 	lastIdx := len(chars.slice)
 	firstIdx := lastIdx - len(runes)