Ignore ANSI sequences when measuring string width (#147)

spinner.go

@@ -496,18 +496,43 @@ func computeNumberOfLinesNeededToPrintString(linePrinted string) int {
 	return computeNumberOfLinesNeededToPrintStringInternal(linePrinted, terminalWidth)
 }
 
-func computeNumberOfLinesNeededToPrintStringInternal(linePrinted string, maxLineWidth int) int {
-	if linePrinted == "" {
-		// empty string will necessarily take one line
-		return 1
+// isAnsiMarker returns if a rune denotes the start of an ANSI sequence
+func isAnsiMarker(r rune) bool {
+	return r == '\x1b'
+}
+
+// isAnsiTerminator returns if a rune denotes the end of an ANSI sequence
+func isAnsiTerminator(r rune) bool {
+	return (r >= 0x40 && r <= 0x5a) || (r == 0x5e) || (r >= 0x60 && r <= 0x7e)
+}
+
+// computeLineWidth returns the displayed width of a line
+func computeLineWidth(line string) int {
+	width := 0
+	ansi := false
+
+	for _, r := range []rune(line) {
+		// increase width only when outside of ANSI escape sequences
+		if ansi || isAnsiMarker(r) {
+			ansi = !isAnsiTerminator(r)
+		} else {
+			width += utf8.RuneLen(r)
+		}
 	}
-	idxOfNewline := strings.Index(linePrinted, "\n")
-	if idxOfNewline < 0 {
-		// we use utf8.RunCountInString() in place of len() because the string contains "complex" unicode chars that
-		// might be represented by multiple individual bytes (typically spinner char)
-		return int(math.Ceil(float64(utf8.RuneCountInString(linePrinted)) / float64(maxLineWidth)))
-	} else {
-		return computeNumberOfLinesNeededToPrintStringInternal(linePrinted[:idxOfNewline], maxLineWidth) +
-			computeNumberOfLinesNeededToPrintStringInternal(linePrinted[idxOfNewline+1:], maxLineWidth)
+
+	return width
+}
+
+func computeNumberOfLinesNeededToPrintStringInternal(linePrinted string, maxLineWidth int) int {
+	lineCount := 0
+	for _, line := range strings.Split(linePrinted, "\n") {
+		lineCount += 1
+
+		lineWidth := computeLineWidth(line)
+		if lineWidth > maxLineWidth {
+			lineCount += int(float64(lineWidth) / float64(maxLineWidth))
+		}
 	}
+
+	return lineCount
 }

spinner_test.go

@@ -281,57 +281,40 @@ func TestWithWriter(t *testing.T) {
 	_ = s
 }
 
-func TestComputeNumberOfLinesNeededToPrintStringInternal_SingleLine(t *testing.T) {
-	line := "Hello world"
-	result := computeNumberOfLinesNeededToPrintStringInternal(line, 50)
-	expectedResult := 1
-	if result != expectedResult {
-		t.Errorf("Line '%s' shoud be printed on '%d' line, got '%d'", line, expectedResult, result)
+func TestComputeNumberOfLinesNeededToPrintStringInternal(t *testing.T) {
+	tests := []struct {
+		description   string
+		expectedCount int
+		printedLine   string
+		maxLineWidth  int
+	}{
+		{"BlankLine", 1, "", 50},
+		{"SingleLine", 1, "Hello world", 50},
+		{"SingleLineANSI", 1, "Hello \x1b[36mworld\x1b[0m", 20},
+		{"MultiLine", 2, "Hello\n world", 50},
+		{"MultiLineANSI", 2, "Hello\n \x1b[1;36mworld\x1b[0m", 20},
+		{"LongString", 2, "Hello world! I am a super long string that will be printed in 2 lines", 50},
+		{"LongStringWithNewlines", 4, "Hello world!\nI am a super long string that will be printed in 2 lines.\nAnother new line", 50},
+		{"NewlineCharAtStart", 2, "\nHello world!", 50},
+		{"NewlineCharAtStartANSI", 2, "\n\x1b[36mHello\x1b[0m world!", 50},
+		{"NewlineCharAtStartANSIFlipped", 2, "\x1b[36m\nHello\x1b[0m world!", 50},
+		{"MultipleNewlineCharAtStart", 4, "\n\n\nHello world!", 50},
+		{"NewlineCharAtEnd", 2, "Hello world!\n", 50},
+		{"NewlineCharAtEndANSI", 2, "Hello \x1b[36mworld!\x1b[0m\n", 50},
+		{"NewlineCharAtEndANSIFlipped", 2, "Hello \x1b[36mworld!\n\x1b[0m", 50},
+		{"StringExactlySizeOfScreen", 1, strings.Repeat("a", 50), 50},
+		{"StringExactlySizeOfScreenANSI", 1, "\x1b[36m" + strings.Repeat("a", 50), 50},
+		{"StringOneGreaterThanSizeOfScreen", 2, strings.Repeat("a", 51), 50},
 	}
-}
-
-func TestComputeNumberOfLinesNeededToPrintStringInternal_MultiLine(t *testing.T) {
-	line := "Hello\n world"
-	result := computeNumberOfLinesNeededToPrintStringInternal(line, 50)
-	expectedResult := 2
-	if result != expectedResult {
-		t.Errorf("Line '%s' shoud be printed on '%d' lines, got '%d'", line, expectedResult, result)
-	}
-}
-
-func TestComputeNumberOfLinesPrinted_LongString(t *testing.T) {
-	line := "Hello world! I am a super long string that will be printed in 2 lines"
-	result := computeNumberOfLinesNeededToPrintStringInternal(line, 50)
-	expectedResult := 2
-	if result != expectedResult {
-		t.Errorf("Line '%s' shoud be printed on '%d' lines, got '%d'", line, expectedResult, result)
-	}
-}
 
-func TestComputeNumberOfLinesNeededToPrintStringInternal_LongStringWithNewlines(t *testing.T) {
-	line := "Hello world!\nI am a super long string that will be printed in 2 lines.\nAnother new line"
-	result := computeNumberOfLinesNeededToPrintStringInternal(line, 50)
-	expectedResult := 4
-	if result != expectedResult {
-		t.Errorf("Line '%s' shoud be printed on '%d' lines, got '%d'", line, expectedResult, result)
-	}
-}
-
-func TestComputeNumberOfLinesNeededToPrintStringInternal_NewlineCharAtTheEnd(t *testing.T) {
-	line := "Hello world!\n"
-	result := computeNumberOfLinesNeededToPrintStringInternal(line, 50)
-	expectedResult := 2
-	if result != expectedResult {
-		t.Errorf("Line '%s' shoud be printed on '%d' lines, got '%d'", line, expectedResult, result)
-	}
-}
-
-func TestComputeNumberOfLinesNeededToPrintStringInternal_StringExactlyTheSizeOfTheScreen(t *testing.T) {
-	line := strings.Repeat("a", 50)
-	result := computeNumberOfLinesNeededToPrintStringInternal(line, 50)
-	expectedResult := 1
-	if result != expectedResult {
-		t.Errorf("Line '%s' shoud be printed on '%d' lines, got '%d'", line, expectedResult, result)
+	for _, test := range tests {
+		result := computeNumberOfLinesNeededToPrintStringInternal(test.printedLine,
+			test.maxLineWidth)
+		if result != test.expectedCount {
+			// Output error, resetting leftover ANSI sequences
+			t.Errorf("%s: Line '%s\x1b[0m' shoud be printed on '%d' line, got '%d'",
+				test.description, test.printedLine, test.expectedCount, result)
+		}
 	}
 }