From 04658fed8eb24c81d1612d89db091e4f6c1be19f Mon Sep 17 00:00:00 2001
From: Kevin Miller <kevin@neomantra.net>
Date: Thu, 11 Apr 2024 07:32:49 -0500
Subject: [PATCH] Add support fo style func to table bubble

---
 table/table.go      | 30 ++++++++++++++++++++++++------
 table/table_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/table/table.go b/table/table.go
index 7bff090..34d5ae8 100644
--- a/table/table.go
+++ b/table/table.go
@@ -14,11 +14,12 @@ import (
 type Model struct {
 	KeyMap KeyMap
 
-	cols   []Column
-	rows   []Row
-	cursor int
-	focus  bool
-	styles Styles
+	cols      []Column
+	rows      []Row
+	cursor    int
+	focus     bool
+	styles    Styles
+	styleFunc StyleFunc
 
 	viewport viewport.Model
 	start    int
@@ -188,6 +189,13 @@ func WithStyles(s Styles) Option {
 	}
 }
 
+// WithStyleFunc sets the table style func which can determine a cell style per column, row, and selected state.
+func WithStyleFunc(f StyleFunc) Option {
+	return func(m *Model) {
+		m.styleFunc = f
+	}
+}
+
 // WithKeyMap sets the key map.
 func WithKeyMap(km KeyMap) Option {
 	return func(m *Model) {
@@ -397,6 +405,9 @@ func (m *Model) FromValues(value, separator string) {
 	m.SetRows(rows)
 }
 
+// StyleFunc is a function that can be used to customize the style of a table cell based on the row and column index.
+type StyleFunc func(row, col int, value string) lipgloss.Style
+
 func (m Model) headersView() string {
 	var s = make([]string, 0, len(m.cols))
 	for _, col := range m.cols {
@@ -416,8 +427,15 @@ func (m *Model) renderRow(rowID int) string {
 		if m.cols[i].Width <= 0 {
 			continue
 		}
+		var cellStyle lipgloss.Style
+		if m.styleFunc != nil {
+			cellStyle = m.styleFunc(rowID, i, value)
+		} else {
+			cellStyle = m.styles.Cell
+		}
+
 		style := lipgloss.NewStyle().Width(m.cols[i].Width).MaxWidth(m.cols[i].Width).Inline(true)
-		renderedCell := m.styles.Cell.Render(style.Render(runewidth.Truncate(value, m.cols[i].Width, "…")))
+		renderedCell := cellStyle.Render(style.Render(runewidth.Truncate(value, m.cols[i].Width, "…")))
 		s = append(s, renderedCell)
 	}
 
diff --git a/table/table_test.go b/table/table_test.go
index d927be0..24f075f 100644
--- a/table/table_test.go
+++ b/table/table_test.go
@@ -1,6 +1,7 @@
 package table
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/charmbracelet/lipgloss"
@@ -108,3 +109,33 @@ func TestRenderRow(t *testing.T) {
 		})
 	}
 }
+func TestRenderRowStyleFunc(t *testing.T) {
+	tests := []struct {
+		name     string
+		table    *Model
+		expected string
+	}{
+		{
+			name: "simple row",
+			table: &Model{
+				rows: []Row{{"Foooooo", "Baaaaar", "Baaaaaz"}},
+				cols: cols,
+				styleFunc: func(row, col int, value string) lipgloss.Style {
+					if strings.HasSuffix(value, "z") {
+						return lipgloss.NewStyle().Transform(strings.ToLower)
+					}
+					return lipgloss.NewStyle().Transform(strings.ToUpper)
+				},
+			},
+			expected: "FOOOOOO   BAAAAAR   baaaaaz   ",
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			row := tc.table.renderRow(0)
+			if row != tc.expected {
+				t.Fatalf("\n\nWant: \n%s\n\nGot:  \n%s\n", tc.expected, row)
+			}
+		})
+	}
+}