Implement an ANSI-aware string truncation function (#22)

ansi/buffer.go

@@ -11,28 +11,77 @@ type Buffer struct {
 	bytes.Buffer
 }
 
-// PrintableRuneWidth returns the width of all printable runes in the buffer.
+// PrintableRuneWidth returns the cell width of all printable runes in the
+// buffer.
 func (w Buffer) PrintableRuneWidth() int {
 	return PrintableRuneWidth(w.String())
 }
 
+// PrintableRuneWidth returns the cell width of the given string.
 func PrintableRuneWidth(s string) int {
 	var n int
 	var ansi bool
 
 	for _, c := range s {
-		if c == '\x1B' {
-			// ANSI escape sequence
-			ansi = true
-		} else if ansi {
-			if (c >= 0x40 && c <= 0x5a) || (c >= 0x61 && c <= 0x7a) {
-				// ANSI sequence terminated
-				ansi = false
-			}
-		} else {
-			n += runewidth.RuneWidth(c)
-		}
+		accPrintableRuneWidth(c, &n, &ansi)
 	}
 
 	return n
 }
+
+// Truncate truncates a string at the given printable cell width, leaving any
+// ansi sequences intact.
+func Truncate(s string, w int) string {
+	return TruncateWithTail(s, w, "")
+}
+
+// TruncateWithTail truncates a string at the given printable cell width,
+// leaving any ansi sequences intact. A tail is then added to the end of the
+// string.
+func TruncateWithTail(s string, w int, tail string) string {
+	if PrintableRuneWidth(s) <= w {
+		return s
+	}
+
+	const ansiReset = "\x1B[0m"
+
+	if tail != "" {
+		tail += ansiReset
+	}
+
+	tw := PrintableRuneWidth(tail)
+	w -= tw
+	if w < 0 {
+		return tail
+	}
+
+	r := []rune(s)
+	ansi := false
+	n := 0
+	i := 0
+
+	for ; i < len(r); i++ {
+		accPrintableRuneWidth(r[i], &n, &ansi)
+		if n > w {
+			break
+		}
+	}
+
+	return string(r[0:i]) + ansiReset + tail
+}
+
+// Used to accumulate the printable rune width while tracking whether we're in
+// an ansi sequence.
+func accPrintableRuneWidth(c rune, n *int, ansi *bool) {
+	if c == '\x1B' {
+		// ANSI escape sequence
+		*ansi = true
+	} else if *ansi {
+		if (c >= 0x40 && c <= 0x5a) || (c >= 0x61 && c <= 0x7a) {
+			// ANSI sequence terminated
+			*ansi = false
+		}
+	} else {
+		*n += runewidth.RuneWidth(c)
+	}
+}

ansi/buffer_test.go

@@ -2,6 +2,7 @@ package ansi
 
 import (
 	"bytes"
+	"fmt"
 	"testing"
 )
 
@@ -29,3 +30,65 @@ func Benchmark_PrintableRuneWidth(b *testing.B) {
 		}
 	})
 }
+
+func Test_Truncate(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		in            string
+		out           string
+		width         int
+		expectedWidth int
+	}{
+		{
+			"\x1B[38;2;249;38;114m你\x1B[7m好\x1B[0m",
+			"\x1B[38;2;249;38;114m你\x1B[7m\x1B[0m",
+			2,
+			2,
+		},
+		{
+			"\x1B[38;2;249;38;114m你\x1B[7m好\x1B[0m",
+			"\x1B[38;2;249;38;114m\x1B[0m",
+			1,
+			0,
+		},
+		{
+			"It’s me!",
+			"It’s me!",
+			10,
+			8,
+		},
+		{
+			"It’s \x1B[7mme!",
+			"It’s \x1B[7m\x1B[0m",
+			5,
+			5,
+		},
+	}
+
+	for i, tt := range tests {
+		t.Run(fmt.Sprintf("truncate-%d", i), func(t *testing.T) {
+			t.Parallel()
+			res := Truncate(tt.in, tt.width)
+			if n := PrintableRuneWidth(res); n != tt.expectedWidth {
+				t.Fatalf("width should be %d, got %d", tt.expectedWidth, n)
+			}
+			if res != tt.out {
+				t.Fatalf("expected '%s' got '%s'\x1B[0m", tt.out, res)
+			}
+		})
+	}
+}
+
+// go test -bench=Benchmark_Truncate -benchmem -count=4
+func Benchmark_Truncate(b *testing.B) {
+	s := "\x1B[38;2;249;38;114m你\x1B[7m好\x1B[0m"
+
+	b.RunParallel(func(pb *testing.PB) {
+		b.ReportAllocs()
+		b.ResetTimer()
+		for pb.Next() {
+			Truncate(s, 2)
+		}
+	})
+}