From de4601232b791a1397d6c645904b001a311a5693 Mon Sep 17 00:00:00 2001
From: Michael Lorant <michael.lorant@nine.com.au>
Date: Sat, 3 Feb 2024 11:04:43 +1100
Subject: [PATCH] Fix truncate of table cells containing ANSI (#256)

---
 go.mod              |  1 +
 go.sum              |  2 ++
 table/table.go      |  6 +++---
 table/table_test.go | 48 +++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index 0123aab..8e7cf2e 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ retract v0.7.0 // v0.7.0 introduces a bug that causes some apps to freeze.
 go 1.17
 
 require (
+	github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d
 	github.com/mattn/go-runewidth v0.0.15
 	github.com/muesli/reflow v0.3.0
 	github.com/muesli/termenv v0.15.2
diff --git a/go.sum b/go.sum
index cba58dc..4a449c2 100644
--- a/go.sum
+++ b/go.sum
@@ -1,3 +1,5 @@
+github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d h1:licZJFw2RwpHMqeKTCYkitsPqHNxTmd4SNR5r94FGM8=
+github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d/go.mod h1:asat636LX7Bqt5lYEZ27JNDcqxfjdBQuJ/MM4CN/Lzo=
 github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
 github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
 github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
diff --git a/table/table.go b/table/table.go
index 953721a..68e41ad 100644
--- a/table/table.go
+++ b/table/table.go
@@ -4,7 +4,7 @@ import (
 	"strings"
 
 	"github.com/charmbracelet/lipgloss"
-	"github.com/mattn/go-runewidth"
+	"github.com/muesli/reflow/truncate"
 )
 
 // StyleFunc is the style function that determines the style of a Cell.
@@ -437,7 +437,7 @@ func (t *Table) constructHeaders() string {
 			MaxHeight(1).
 			Width(t.widths[i]).
 			MaxWidth(t.widths[i]).
-			Render(runewidth.Truncate(header, t.widths[i], "…")))
+			Render(truncate.StringWithTail(header, uint(t.widths[i]), "…")))
 		if i < len(t.headers)-1 && t.borderColumn {
 			s.WriteString(t.borderStyle.Render(t.border.Left))
 		}
@@ -488,7 +488,7 @@ func (t *Table) constructRow(index int) string {
 			MaxHeight(height).
 			Width(t.widths[c]).
 			MaxWidth(t.widths[c]).
-			Render(runewidth.Truncate(cell, t.widths[c]*height, "…")))
+			Render(truncate.StringWithTail(cell, uint(t.widths[c]*height), "…")))
 
 		if c < t.data.Columns()-1 && t.borderColumn {
 			cells = append(cells, left)
diff --git a/table/table_test.go b/table/table_test.go
index 4e5ccad..31f3010 100644
--- a/table/table_test.go
+++ b/table/table_test.go
@@ -3,7 +3,9 @@ package table
 import (
 	"strings"
 	"testing"
+	"unicode"
 
+	"github.com/acarl005/stripansi"
 	"github.com/charmbracelet/lipgloss"
 )
 
@@ -941,6 +943,52 @@ func TestFilterInverse(t *testing.T) {
 	}
 }
 
+func TestTableANSI(t *testing.T) {
+	const code = "\x1b[31mC\x1b[0m\x1b[32mo\x1b[0m\x1b[34md\x1b[0m\x1b[33me\x1b[0m"
+
+	rows := [][]string{
+		{"Apple", "Red", "\x1b[31m31\x1b[0m"},
+		{"Lime", "Green", "\x1b[32m32\x1b[0m"},
+		{"Banana", "Yellow", "\x1b[33m33\x1b[0m"},
+		{"Blueberry", "Blue", "\x1b[34m34\x1b[0m"},
+	}
+
+	table := New().
+		Width(29).
+		StyleFunc(TableStyle).
+		Border(lipgloss.NormalBorder()).
+		Headers("Fruit", "Color", code).
+		Rows(rows...)
+
+	expected := strings.TrimSpace(`
+┌───────────┬────────┬──────┐
+│   Fruit   │ Color  │ Code │
+├───────────┼────────┼──────┤
+│ Apple     │ Red    │ 31   │
+│ Lime      │ Green  │ 32   │
+│ Banana    │ Yellow │ 33   │
+│ Blueberry │ Blue   │ 34   │
+└───────────┴────────┴──────┘
+`)
+
+	if stripString(table.String()) != expected {
+		t.Fatalf("expected:\n\n%s\n\ngot:\n\n%s", expected, stripString(table.String()))
+	}
+}
+
 func debug(s string) string {
 	return strings.ReplaceAll(s, " ", ".")
 }
+
+func stripString(str string) string {
+	s := stripansi.Strip(str)
+	ss := strings.Split(s, "\n")
+
+	var lines []string
+	for _, l := range ss {
+		trim := strings.TrimRightFunc(l, unicode.IsSpace)
+		lines = append(lines, trim)
+	}
+
+	return strings.Join(lines, "\n")
+}