Remove ANSI formatting before measuring string width (#462)

Author: zimeg

go.sum

@@ -21,7 +21,6 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
-golang.org/x/sys v0.0.0-20201119102817-f84b799fce68 h1:nxC68pudNYkKU6jWhgrqdreuFiOQWj1Fs7T3VrH4Pjw=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20220422013727-9388b58f7150 h1:xHms4gcpe1YE7A3yIllJXP16CMAGuqwO2lX1mTyyRRc=
 golang.org/x/sys v0.0.0-20220422013727-9388b58f7150/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

terminal/runereader.go

@@ -377,19 +377,42 @@ func (rr *RuneReader) ReadLineWithDefault(mask rune, d []rune, onRunes ...OnRune
 	}
 }
 
+// runeWidth returns the number of columns spanned by a rune when printed to the terminal
 func runeWidth(r rune) int {
 	switch width.LookupRune(r).Kind() {
 	case width.EastAsianWide, width.EastAsianFullwidth:
 		return 2
 	}
+
+	if !unicode.IsPrint(r) {
+		return 0
+	}
 	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)
+}
+
+// StringWidth returns the visible width of a string when printed to the terminal
 func StringWidth(str string) int {
 	w := 0
+	ansi := false
+
 	rs := []rune(str)
 	for _, r := range rs {
-		w += runeWidth(r)
+		// increase width only when outside of ANSI escape sequences
+		if ansi || isAnsiMarker(r) {
+			ansi = !isAnsiTerminator(r)
+		} else {
+			w += runeWidth(r)
+		}
 	}
 	return w
 }

terminal/runereader_test.go

@@ -0,0 +1,66 @@
+package terminal
+
+import (
+	"testing"
+)
+
+func TestRuneWidthInvisible(t *testing.T) {
+	var example rune = '⁣'
+	expected := 0
+	actual := runeWidth(example)
+	if actual != expected {
+		t.Errorf("Expected '%c' to have width %d, found %d", example, expected, actual)
+	}
+}
+
+func TestRuneWidthNormal(t *testing.T) {
+	var example rune = 'a'
+	expected := 1
+	actual := runeWidth(example)
+	if actual != expected {
+		t.Errorf("Expected '%c' to have width %d, found %d", example, expected, actual)
+	}
+}
+
+func TestRuneWidthWide(t *testing.T) {
+	var example rune = '错'
+	expected := 2
+	actual := runeWidth(example)
+	if actual != expected {
+		t.Errorf("Expected '%c' to have width %d, found %d", example, expected, actual)
+	}
+}
+
+func TestStringWidthEmpty(t *testing.T) {
+	example := ""
+	expected := 0
+	actual := StringWidth(example)
+	if actual != expected {
+		t.Errorf("Expected '%s' to have width %d, found %d", example, expected, actual)
+	}
+}
+
+func TestStringWidthNormal(t *testing.T) {
+	example := "Green"
+	expected := 5
+	actual := StringWidth(example)
+	if actual != expected {
+		t.Errorf("Expected '%s' to have width %d, found %d", example, expected, actual)
+	}
+}
+
+func TestStringWidthFormat(t *testing.T) {
+	example := "\033[31mRed\033[0m"
+	expected := 3
+	actual := StringWidth(example)
+	if actual != expected {
+		t.Errorf("Expected '%s' to have width %d, found %d", example, expected, actual)
+	}
+
+	example = "\033[1;34mbold\033[21mblue\033[0m"
+	expected = 8
+	actual = StringWidth(example)
+	if actual != expected {
+		t.Errorf("Expected '%s' to have width %d, found %d", example, expected, actual)
+	}
+}