From b54ce3849e4770c86edea7108c23e8e2c89fea2b Mon Sep 17 00:00:00 2001
From: Ethan Zimbelman <ethan.zimbelman@me.com>
Date: Mon, 7 Nov 2022 08:11:34 -0800
Subject: [PATCH] Remove ANSI formatting before measuring string width (#462)

---
 go.sum                      |  1 -
 terminal/runereader.go      | 25 +++++++++++++-
 terminal/runereader_test.go | 66 +++++++++++++++++++++++++++++++++++++
 3 files changed, 90 insertions(+), 2 deletions(-)
 create mode 100644 terminal/runereader_test.go

diff --git a/go.sum b/go.sum
index e83f9839..998b52d5 100644
--- a/go.sum
+++ b/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=
diff --git a/terminal/runereader.go b/terminal/runereader.go
index c6997597..eab41f7f 100644
--- a/terminal/runereader.go
+++ b/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
 }
diff --git a/terminal/runereader_test.go b/terminal/runereader_test.go
new file mode 100644
index 00000000..fe8fc881
--- /dev/null
+++ b/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)
+	}
+}