text: fix parsing escape sequences while wrapping; fixes #330 (#333)

Author: jedib0t

text/escape_sequences.go

@@ -0,0 +1,96 @@
+package text
+
+import (
+	"fmt"
+	"sort"
+	"strconv"
+	"strings"
+)
+
+type escSeqParser struct {
+	openSeq map[int]bool
+}
+
+func (s *escSeqParser) Codes() []int {
+	codes := make([]int, 0)
+	for code, val := range s.openSeq {
+		if val {
+			codes = append(codes, code)
+		}
+	}
+	sort.Ints(codes)
+	return codes
+}
+
+func (s *escSeqParser) Extract(str string) string {
+	escapeSeq, inEscSeq := "", false
+	for _, char := range str {
+		if char == EscapeStartRune {
+			inEscSeq = true
+			escapeSeq = ""
+		}
+		if inEscSeq {
+			escapeSeq += string(char)
+		}
+		if char == EscapeStopRune {
+			inEscSeq = false
+			s.Parse(escapeSeq)
+		}
+	}
+	return s.Sequence()
+}
+
+func (s *escSeqParser) IsOpen() bool {
+	return len(s.openSeq) > 0
+}
+
+func (s *escSeqParser) Sequence() string {
+	out := strings.Builder{}
+	if s.IsOpen() {
+		out.WriteString(EscapeStart)
+		for idx, code := range s.Codes() {
+			if idx > 0 {
+				out.WriteRune(';')
+			}
+			out.WriteString(fmt.Sprint(code))
+		}
+		out.WriteString(EscapeStop)
+	}
+
+	return out.String()
+}
+
+func (s *escSeqParser) Parse(seq string) {
+	if s.openSeq == nil {
+		s.openSeq = make(map[int]bool)
+	}
+
+	seq = strings.Replace(seq, EscapeStart, "", 1)
+	seq = strings.Replace(seq, EscapeStop, "", 1)
+	codes := strings.Split(seq, ";")
+	for _, code := range codes {
+		code = strings.TrimSpace(code)
+		if codeNum, err := strconv.Atoi(code); err == nil {
+			switch codeNum {
+			case 0: // reset
+				s.openSeq = make(map[int]bool) // clear everything
+			case 22: // reset intensity
+				delete(s.openSeq, 1) // remove bold
+				delete(s.openSeq, 2) // remove faint
+			case 23: // not italic
+				delete(s.openSeq, 3) // remove italic
+			case 24: // not underlined
+				delete(s.openSeq, 4) // remove underline
+			case 25: // not blinking
+				delete(s.openSeq, 5) // remove slow blink
+				delete(s.openSeq, 6) // remove rapid blink
+			case 27: // not reversed
+				delete(s.openSeq, 7) // remove reverse
+			case 29: // not crossed-out
+				delete(s.openSeq, 9) // remove crossed-out
+			default:
+				s.openSeq[codeNum] = true
+			}
+		}
+	}
+}

text/escape_sequences_test.go

@@ -0,0 +1,41 @@
+package text
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func Test_escSeqParser(t *testing.T) {
+	t.Run("extract", func(t *testing.T) {
+		es := escSeqParser{}
+
+		assert.Equal(t, "\x1b[1;3;4;5;7;9;91m", es.Extract("\x1b[91m\x1b[1m\x1b[3m\x1b[4m\x1b[5m\x1b[7m\x1b[9m Spicy"))
+		assert.Equal(t, "\x1b[3;4;5;7;9;91m", es.Extract("\x1b[22m No Bold"))
+		assert.Equal(t, "\x1b[4;5;7;9;91m", es.Extract("\x1b[23m No Italic"))
+		assert.Equal(t, "\x1b[5;7;9;91m", es.Extract("\x1b[24m No Underline"))
+		assert.Equal(t, "\x1b[7;9;91m", es.Extract("\x1b[25m No Blink"))
+		assert.Equal(t, "\x1b[9;91m", es.Extract("\x1b[27m No Reverse"))
+		assert.Equal(t, "\x1b[91m", es.Extract("\x1b[29m No Crossed-Out"))
+		assert.Equal(t, "", es.Extract("\x1b[0m Resetted"))
+	})
+
+	t.Run("parse", func(t *testing.T) {
+		es := escSeqParser{}
+
+		es.Parse("\x1b[91m") // color
+		es.Parse("\x1b[1m")  // bold
+		assert.Len(t, es.Codes(), 2)
+		assert.True(t, es.IsOpen())
+		assert.Equal(t, "\x1b[1;91m", es.Sequence())
+
+		es.Parse("\x1b[22m") // un-bold
+		assert.Len(t, es.Codes(), 1)
+		assert.True(t, es.IsOpen())
+		assert.Equal(t, "\x1b[91m", es.Sequence())
+
+		es.Parse("\x1b[0m") // reset
+		assert.Empty(t, es.Codes())
+		assert.False(t, es.IsOpen())
+		assert.Empty(t, es.Sequence())
+	})
+}

text/wrap.go

@@ -69,32 +69,21 @@ func WrapText(str string, wrapLen int) string {
 	if wrapLen <= 0 {
 		return ""
 	}
-
-	var out strings.Builder
+	str = strings.Replace(str, "\t", "    ", -1)
 	sLen := utf8.RuneCountInString(str)
-	out.Grow(sLen + (sLen / wrapLen))
-	lineIdx, isEscSeq, lastEscSeq := 0, false, ""
-	for _, char := range str {
-		if char == EscapeStartRune {
-			isEscSeq = true
-			lastEscSeq = ""
-		}
-		if isEscSeq {
-			lastEscSeq += string(char)
-		}
-
-		appendChar(char, wrapLen, &lineIdx, isEscSeq, lastEscSeq, &out)
+	if sLen <= wrapLen {
+		return str
+	}
 
-		if isEscSeq && char == EscapeStopRune {
-			isEscSeq = false
-		}
-		if lastEscSeq == EscapeReset {
-			lastEscSeq = ""
+	out := &strings.Builder{}
+	out.Grow(sLen + (sLen / wrapLen))
+	for idx, line := range strings.Split(str, "\n") {
+		if idx > 0 {
+			out.WriteString("\n")
 		}
+		wrapHard(line, wrapLen, out)
 	}
-	if lastEscSeq != "" && lastEscSeq != EscapeReset {
-		out.WriteString(EscapeReset)
-	}
+
 	return out.String()
 }
 
@@ -149,26 +138,6 @@ func appendWord(word string, lineIdx *int, lastSeenEscSeq string, wrapLen int, o
 	}
 }
 
-func extractOpenEscapeSeq(str string) string {
-	escapeSeq, inEscSeq := "", false
-	for _, char := range str {
-		if char == EscapeStartRune {
-			inEscSeq = true
-			escapeSeq = ""
-		}
-		if inEscSeq {
-			escapeSeq += string(char)
-		}
-		if char == EscapeStopRune {
-			inEscSeq = false
-		}
-	}
-	if escapeSeq == EscapeReset {
-		escapeSeq = ""
-	}
-	return escapeSeq
-}
-
 func terminateLine(wrapLen int, lineLen *int, lastSeenEscSeq string, out *strings.Builder) {
 	if *lineLen < wrapLen {
 		out.WriteString(strings.Repeat(" ", wrapLen-*lineLen))
@@ -189,12 +158,12 @@ func terminateOutput(lastSeenEscSeq string, out *strings.Builder) {
 }
 
 func wrapHard(paragraph string, wrapLen int, out *strings.Builder) {
+	esp := escSeqParser{}
 	lineLen, lastSeenEscSeq := 0, ""
 	words := strings.Fields(paragraph)
 	for wordIdx, word := range words {
-		escSeq := extractOpenEscapeSeq(word)
-		if escSeq != "" {
-			lastSeenEscSeq = escSeq
+		if openEscSeq := esp.Extract(word); openEscSeq != "" {
+			lastSeenEscSeq = openEscSeq
 		}
 		if lineLen > 0 {
 			out.WriteRune(' ')
@@ -218,12 +187,12 @@ func wrapHard(paragraph string, wrapLen int, out *strings.Builder) {
 }
 
 func wrapSoft(paragraph string, wrapLen int, out *strings.Builder) {
+	esp := escSeqParser{}
 	lineLen, lastSeenEscSeq := 0, ""
 	words := strings.Fields(paragraph)
 	for wordIdx, word := range words {
-		escSeq := extractOpenEscapeSeq(word)
-		if escSeq != "" {
-			lastSeenEscSeq = escSeq
+		if openEscSeq := esp.Extract(word); openEscSeq != "" {
+			lastSeenEscSeq = openEscSeq
 		}
 
 		spacing, spacingLen := wrapSoftSpacing(lineLen)

text/wrap_test.go

@@ -50,6 +50,15 @@ func TestWrapHard(t *testing.T) {
 
 	complexIn := "+---+------+-------+------+\n| 1 | Arya | Stark | 3000 |\n+---+------+-------+------+"
 	assert.Equal(t, complexIn, WrapHard(complexIn, 27))
+
+	// colored text with nested escape codes: "{red}{bold}...{un-bold}...{reset}"
+	textUnBold := "\x1b[91m\x1b[1mBold Title\x1b[22m Regular Red Text\x1b[0m"
+	expectedUnBold := "\x1b[91m\x1b[1mBold Title\x1b[22m Regular Red \x1b[0m\n\x1b[91mText\x1b[0m"
+	assert.Equal(t, expectedUnBold, WrapHard(textUnBold, 23))
+}
+
+func TestFoo(t *testing.T) {
+	assert.Equal(t, "\x1b[33mJon\x1b[0m\n\x1b[33mSno\x1b[0m\n\x1b[33mw\x1b[0m", WrapHard("\x1b[33mJon Snow\x1b[0m", 3))
 }
 
 func ExampleWrapSoft() {
@@ -100,6 +109,11 @@ func TestWrapSoft(t *testing.T) {
 
 	assert.Equal(t, "\x1b[33mJon \x1b[0m\n\x1b[33mSnow\x1b[0m", WrapSoft("\x1b[33mJon Snow\x1b[0m", 4))
 	assert.Equal(t, "\x1b[33mJon \x1b[0m\n\x1b[33mSnow\x1b[0m\n\x1b[33m???\x1b[0m", WrapSoft("\x1b[33mJon Snow???\x1b[0m", 4))
+
+	// colored text with nested escape codes: "{red}{bold}...{un-bold}...{reset}"
+	textUnBold := "\x1b[91m\x1b[1mBold Title\x1b[22m Regular Red Text\x1b[0m"
+	expectedUnBold := "\x1b[91m\x1b[1mBold Title\x1b[22m Regular Red \x1b[0m\n\x1b[91mText\x1b[0m"
+	assert.Equal(t, expectedUnBold, WrapSoft(textUnBold, 23))
 }
 
 func ExampleWrapText() {
@@ -138,10 +152,15 @@ func TestWrapText(t *testing.T) {
 	assert.Equal(t, "Jon\nSno\nw\n", WrapText("Jon\nSnow\n", 3))
 	assert.Equal(t, "\x1b[33mJon\x1b[0m\nSno\nw", WrapText("\x1b[33mJon\x1b[0m\nSnow", 3))
 	assert.Equal(t, "\x1b[33mJon\x1b[0m\nSno\nw\n", WrapText("\x1b[33mJon\x1b[0m\nSnow\n", 3))
-	assert.Equal(t, "\x1b[33mJon\x1b[0m\n\x1b[33m Sn\x1b[0m\n\x1b[33mow\x1b[0m", WrapText("\x1b[33mJon Snow\x1b[0m", 3))
-	assert.Equal(t, "\x1b[33mJon\x1b[0m\n\x1b[33m Sn\x1b[0m\n\x1b[33mow\x1b[0m\n\x1b[33m\x1b[0m", WrapText("\x1b[33mJon Snow\n", 3))
-	assert.Equal(t, "\x1b[33mJon\x1b[0m\n\x1b[33m Sn\x1b[0m\n\x1b[33mow\x1b[0m\n\x1b[33m\x1b[0m", WrapText("\x1b[33mJon Snow\n\x1b[0m", 3))
+	assert.Equal(t, "\x1b[33mJon\x1b[0m\n\x1b[33mSno\x1b[0m\n\x1b[33mw\x1b[0m", WrapText("\x1b[33mJon Snow\x1b[0m", 3))
+	assert.Equal(t, "\x1b[33mJon\x1b[0m\n\x1b[33mSno\x1b[0m\n\x1b[33mw\x1b[0m\n", WrapText("\x1b[33mJon Snow\n", 3))
+	assert.Equal(t, "\x1b[33mJon\x1b[0m\n\x1b[33mSno\x1b[0m\n\x1b[33mw\x1b[0m\n\x1b[0m", WrapText("\x1b[33mJon Snow\n\x1b[0m", 3))
 
 	complexIn := "+---+------+-------+------+\n| 1 | Arya | Stark | 3000 |\n+---+------+-------+------+"
 	assert.Equal(t, complexIn, WrapText(complexIn, 27))
+
+	// colored text with nested escape codes: "{red}{bold}...{un-bold}...{reset}"
+	textUnBold := "\x1b[91m\x1b[1mBold Title\x1b[22m Regular Red Text\x1b[0m"
+	expectedUnBold := "\x1b[91m\x1b[1mBold Title\x1b[22m Regular Red \x1b[0m\n\x1b[91mText\x1b[0m"
+	assert.Equal(t, expectedUnBold, WrapText(textUnBold, 23))
 }