termui: properly handle color sequence code even inside a word

2 files changed, 191 insertions, 34 deletions

diff --git a/util/text.go b/util/text.go
index cc652cc6..3b2b0020 100644
--- a/util/text.go
+++ b/util/text.go
@@ -2,7 +2,6 @@ package util
 
 import (
 	"bytes"
-	"regexp"
 	"strings"
 )
 
@@ -28,10 +27,14 @@ func WordWrap(text string, lineWidth int) (string, int) {
 	return wrapped, lines
 }
 
+// Wrap a text for an exact line size
+// Handle properly terminal color escape code
 func TextWrap(text string, lineWidth int) (string, int) {
 	return TextWrapPadded(text, lineWidth, 0)
 }
 
+// Wrap a text for an exact line size with a left padding
+// Handle properly terminal color escape code
 func TextWrapPadded(text string, lineWidth int, leftPad int) (string, int) {
 	var textBuffer bytes.Buffer
 	var lineBuffer bytes.Buffer
@@ -42,8 +45,6 @@ func TextWrapPadded(text string, lineWidth int, leftPad int) (string, int) {
 	// tabs are formatted as 4 spaces
 	text = strings.Replace(text, "\t", "    ", 4)
 
-	re := regexp.MustCompile(`(\x1b\[\d+m)?([^\x1b]*)(\x1b\[\d+m)?`)
-
 	for _, line := range strings.Split(text, "\n") {
 		spaceLeft := lineWidth - leftPad
 
@@ -55,56 +56,62 @@ func TextWrapPadded(text string, lineWidth int, leftPad int) (string, int) {
 		firstWord := true
 
 		for _, word := range strings.Split(line, " ") {
-			prefix := ""
-			suffix := ""
-
-			matches := re.FindStringSubmatch(word)
-			if matches != nil && (matches[1] != "" || matches[3] != "") {
-				// we have a color escape sequence
-				prefix = matches[1]
-				word = matches[2]
-				suffix = matches[3]
-			}
+			wordLength := wordLen(word)
+
+			if !firstWord {
+				lineBuffer.WriteString(" ")
+				spaceLeft -= 1
 
-			if spaceLeft > len(word) {
-				if !firstWord {
-					lineBuffer.WriteString(" ")
-					spaceLeft -= 1
+				if spaceLeft <= 0 {
+					textBuffer.WriteString(pad + strings.TrimRight(lineBuffer.String(), " "))
+					textBuffer.WriteString("\n")
+					lineBuffer.Reset()
+					spaceLeft = lineWidth - leftPad
+					nbLine++
+					firstLine = false
 				}
-				lineBuffer.WriteString(prefix + word + suffix)
-				spaceLeft -= len(word)
+			}
+
+			// Word fit in the current line
+			if spaceLeft >= wordLength {
+				lineBuffer.WriteString(word)
+				spaceLeft -= wordLength
 				firstWord = false
 			} else {
-				if len(word) > lineWidth {
-					for len(word) > 0 {
-						l := minInt(spaceLeft, len(word))
-						part := prefix + word[:l]
-						prefix = ""
-						word = word[l:]
+				// Break a word longer than a line
+				if wordLength > lineWidth {
+					for wordLength > 0 && len(word) > 0 {
+						l := minInt(spaceLeft, wordLength)
+						part, leftover := splitWord(word, l)
+						word = leftover
+						wordLength = wordLen(word)
 
 						lineBuffer.WriteString(part)
 						textBuffer.WriteString(pad)
 						textBuffer.Write(lineBuffer.Bytes())
 						lineBuffer.Reset()
 
-						if len(word) > 0 {
+						spaceLeft -= l
+
+						if spaceLeft <= 0 {
 							textBuffer.WriteString("\n")
 							nbLine++
+							spaceLeft = lineWidth - leftPad
 						}
-
-						spaceLeft = lineWidth - leftPad
 					}
 				} else {
+					// Normal break
 					textBuffer.WriteString(pad + strings.TrimRight(lineBuffer.String(), " "))
 					textBuffer.WriteString("\n")
 					lineBuffer.Reset()
-					lineBuffer.WriteString(prefix + word + suffix)
+					lineBuffer.WriteString(word)
 					firstWord = false
-					spaceLeft = lineWidth - len(word)
+					spaceLeft = lineWidth - wordLength
 					nbLine++
 				}
 			}
 		}
+
 		textBuffer.WriteString(pad + strings.TrimRight(lineBuffer.String(), " "))
 		lineBuffer.Reset()
 		firstLine = false
@@ -113,6 +120,60 @@ func TextWrapPadded(text string, lineWidth int, leftPad int) (string, int) {
 	return textBuffer.String(), nbLine
 }
 
+func wordLen(word string) int {
+	length := 0
+	escape := false
+
+	for _, char := range word {
+		if char == '\x1b' {
+			escape = true
+		}
+
+		if !escape {
+			length++
+		}
+
+		if char == 'm' {
+			escape = false
+		}
+	}
+
+	return length
+}
+
+func splitWord(word string, length int) (string, string) {
+	result := ""
+	added := 0
+	escape := false
+
+	if length == 0 {
+		return "", word
+	}
+
+	for _, char := range word {
+		if char == '\x1b' {
+			escape = true
+		}
+
+		result += string(char)
+
+		if !escape {
+			added++
+			if added == length {
+				break
+			}
+		}
+
+		if char == 'm' {
+			escape = false
+		}
+	}
+
+	leftover := word[len(result):]
+
+	return result, leftover
+}
+
 func minInt(a, b int) int {
 	if a > b {
 		return b
diff --git a/util/text_test.go b/util/text_test.go
index 09b6b4d6..a368b329 100644
--- a/util/text_test.go
+++ b/util/text_test.go
@@ -43,7 +43,7 @@ func TestTextWrap(t *testing.T) {
 		// A tab counts as 4 characters.
 		{
 			"foo\nb\t r\n baz",
-			"foo\nb\n r\n baz",
+			"foo\nb\n  r\n baz",
 			4,
 		},
 		// Trailing whitespace is removed after used for wrapping.
@@ -71,10 +71,22 @@ func TestTextWrap(t *testing.T) {
 			"foo\n\x1b[31mbar\x1b[0m\nbaz",
 			4,
 		},
+		// Handle words with colors sequence inside the word
+		{
+			"foo b\x1b[31mbar\x1b[0mr baz",
+			"foo\nb\x1b[31mbar\n\x1b[0mr\nbaz",
+			4,
+		},
+		// Break words with colors sequence inside the word
+		{
+			"foo bb\x1b[31mbar\x1b[0mr baz",
+			"foo\nbb\x1b[31mba\nr\x1b[0mr\nbaz",
+			4,
+		},
 		// Complete example:
 		{
 			" This is a list: \n\n\t* foo\n\t* bar\n\n\n\t* baz  \nBAM    ",
-			" This\nis a\nlist:\n\n    *\nfoo\n    *\nbar\n\n\n    *\nbaz\nBAM\n",
+			" This\nis a\nlist:\n\n\n    *\nfoo\n    *\nbar\n\n\n    *\nbaz\nBAM\n",
 			6,
 		},
 	}
@@ -88,8 +100,92 @@ func TestTextWrap(t *testing.T) {
 
 		expected := len(strings.Split(tc.Output, "\n"))
 		if expected != lines {
-			t.Fatalf("Nb lines mismatch\nExpected:%d\nActual:%d",
-				expected, lines)
+			t.Fatalf("Case %d Nb lines mismatch\nExpected:%d\nActual:%d",
+				i, expected, lines)
+		}
+	}
+}
+
+func TestWordLen(t *testing.T) {
+	cases := []struct {
+		Input  string
+		Length int
+	}{
+		// A simple word
+		{
+			"foo",
+			3,
+		},
+		// A simple word with colors
+		{
+			"\x1b[31mbar\x1b[0m",
+			3,
+		},
+		// Handle prefix and suffix properly
+		{
+			"foo\x1b[31mfoobarHoy\x1b[0mbaaar",
+			17,
+		},
+	}
+
+	for i, tc := range cases {
+		l := wordLen(tc.Input)
+		if l != tc.Length {
+			t.Fatalf("Case %d Input:\n\n`%s`\n\nExpected Output:\n\n`%d`\n\nActual Output:\n\n`%d`",
+				i, tc.Input, tc.Length, l)
+		}
+	}
+}
+
+func TestSplitWord(t *testing.T) {
+	cases := []struct {
+		Input            string
+		Length           int
+		Result, Leftover string
+	}{
+		// A simple word passes through.
+		{
+			"foo",
+			4,
+			"foo", "",
+		},
+		// Cut at the right place
+		{
+			"foobarHoy",
+			4,
+			"foob", "arHoy",
+		},
+		// A simple word passes through with colors
+		{
+			"\x1b[31mbar\x1b[0m",
+			4,
+			"\x1b[31mbar\x1b[0m", "",
+		},
+		// Cut at the right place with colors
+		{
+			"\x1b[31mfoobarHoy\x1b[0m",
+			4,
+			"\x1b[31mfoob", "arHoy\x1b[0m",
+		},
+		// Handle prefix and suffix properly
+		{
+			"foo\x1b[31mfoobarHoy\x1b[0mbaaar",
+			4,
+			"foo\x1b[31mf", "oobarHoy\x1b[0mbaaar",
+		},
+		// Cut properly with length = 0
+		{
+			"foo",
+			0,
+			"", "foo",
+		},
+	}
+
+	for i, tc := range cases {
+		result, leftover := splitWord(tc.Input, tc.Length)
+		if result != tc.Result || leftover != tc.Leftover {
+			t.Fatalf("Case %d Input:\n\n`%s`\n\nExpected Output:\n\n`%s` - `%s`\n\nActual Output:\n\n`%s` - `%s`",
+				i, tc.Input, tc.Result, tc.Leftover, result, leftover)
 		}
 	}
 }