Add cell style func (#203)

* Add cell style func

* Ignore the lint on funlen for now

* Proper nolint, whoops

Author: Evertras

examples/pokemon/main.go

@@ -35,10 +35,24 @@ type Model struct {
 	pokeTable table.Model
 }
 
-func makeRow(name, element, colorStr string, numConversations int, positiveSentiment, negativeSentiment float32) table.Row {
+func makeRow(name, element string, numConversations int, positiveSentiment, negativeSentiment float32) table.Row {
+	elementStyleFunc := func(input table.StyledCellFuncInput) lipgloss.Style {
+		switch input.Data.(string) {
+		case "Fire":
+			return lipgloss.NewStyle().Foreground(lipgloss.Color(colorFire))
+		case "Water":
+			return lipgloss.NewStyle().Foreground(lipgloss.Color(colorWater))
+		case "Plant":
+			return lipgloss.NewStyle().Foreground(lipgloss.Color(colorPlant))
+		case "Electric":
+			return lipgloss.NewStyle().Foreground(lipgloss.Color(colorElectric))
+		default:
+			return lipgloss.NewStyle().Foreground(lipgloss.Color(colorNormal))
+		}
+	}
 	return table.NewRow(table.RowData{
 		columnKeyName:              name,
-		columnKeyElement:           table.NewStyledCell(element, lipgloss.NewStyle().Foreground(lipgloss.Color(colorStr))),
+		columnKeyElement:           table.NewStyledCellWithStyleFunc(element, elementStyleFunc),
 		columnKeyConversations:     numConversations,
 		columnKeyPositiveSentiment: positiveSentiment,
 		columnKeyNegativeSentiment: negativeSentiment,
@@ -58,13 +72,13 @@ func NewModel() Model {
 				WithStyle(lipgloss.NewStyle().Foreground(lipgloss.Color("#c88"))).
 				WithFormatString("%.1f%%"),
 		}).WithRows([]table.Row{
-			makeRow("Pikachu", "Electric", colorElectric, 2300648, 21.9, 8.54),
-			makeRow("Eevee", "Normal", colorNormal, 636373, 26.4, 7.37),
-			makeRow("Bulbasaur", "Plant", colorPlant, 352190, 25.7, 9.02),
-			makeRow("Squirtle", "Water", colorWater, 241259, 25.6, 5.96),
-			makeRow("Blastoise", "Water", colorWater, 162794, 19.5, 6.04),
-			makeRow("Charmander", "Fire", colorFire, 265760, 31.2, 5.25),
-			makeRow("Charizard", "Fire", colorFire, 567763, 25.6, 7.56),
+			makeRow("Pikachu", "Electric", 2300648, 21.9, 8.54),
+			makeRow("Eevee", "Normal", 636373, 26.4, 7.37),
+			makeRow("Bulbasaur", "Plant", 352190, 25.7, 9.02),
+			makeRow("Squirtle", "Water", 241259, 25.6, 5.96),
+			makeRow("Blastoise", "Water", 162794, 19.5, 6.04),
+			makeRow("Charmander", "Fire", 265760, 31.2, 5.25),
+			makeRow("Charizard", "Fire", 567763, 25.6, 7.56),
 		}).
 			BorderRounded().
 			WithBaseStyle(styleBase).

table/cell.go

@@ -8,12 +8,47 @@ import "github.com/charmbracelet/lipgloss"
 // limited to colors, font style, and alignments - spacing style such as margin
 // will break the table format.
 type StyledCell struct {
-	Data  any
+	// Data is the content of the cell.
+	Data any
+
+	// Style is the specific style to apply. This is ignored if StyleFunc is not nil.
 	Style lipgloss.Style
+
+	// StyleFunc is a function that takes the row/column of the cell and
+	// returns a lipgloss.Style allowing for dynamic styling based on the cell's
+	// content or position. Overrides Style if set.
+	StyleFunc StyledCellFunc
+}
+
+// StyledCellFuncInput is the input to the StyledCellFunc. Sent as a struct
+// to allow for future additions without breaking changes.
+type StyledCellFuncInput struct {
+	// Data is the data in the cell.
+	Data any
+
+	// Column is the column that the cell belongs to.
+	Column Column
 }
 
+// StyledCellFunc is a function that takes various information about the cell and
+// returns a lipgloss.Style allowing for easier dynamic styling based on the cell's
+// content or position.
+type StyledCellFunc = func(input StyledCellFuncInput) lipgloss.Style
+
 // NewStyledCell creates an entry that can be set in the row data and show as
 // styled with the given style.
 func NewStyledCell(data any, style lipgloss.Style) StyledCell {
-	return StyledCell{data, style}
+	return StyledCell{
+		Data:  data,
+		Style: style,
+	}
+}
+
+// NewStyledCellWithStyleFunc creates an entry that can be set in the row data and show as
+// styled with the given style function.
+func NewStyledCellWithStyleFunc(data any, styleFunc StyledCellFunc) StyledCell {
+	return StyledCell{
+		Data:      data,
+		StyleFunc: styleFunc,
+	}
 }

table/row.go

@@ -53,7 +53,7 @@ func (r Row) WithStyle(style lipgloss.Style) Row {
 	return r
 }
 
-//nolint:cyclop // This has many ifs, but they're short
+//nolint:cyclop,funlen // Breaking this up will be more complicated than it's worth for now
 func (m Model) renderRowColumnData(row Row, column Column, rowStyle lipgloss.Style, borderStyle lipgloss.Style) string {
 	cellStyle := rowStyle.Copy().Inherit(column.style).Inherit(m.baseStyle)
 
@@ -91,7 +91,15 @@ func (m Model) renderRowColumnData(row Row, column Column, rowStyle lipgloss.Sty
 		switch entry := data.(type) {
 		case StyledCell:
 			str = fmt.Sprintf(fmtString, entry.Data)
-			cellStyle = entry.Style.Copy().Inherit(cellStyle)
+
+			if entry.StyleFunc != nil {
+				cellStyle = entry.StyleFunc(StyledCellFuncInput{
+					Column: column,
+					Data:   entry.Data,
+				}).Copy().Inherit(cellStyle)
+			} else {
+				cellStyle = entry.Style.Copy().Inherit(cellStyle)
+			}
 		default:
 			str = fmt.Sprintf(fmtString, entry)
 		}

table/view_test.go

@@ -798,7 +798,55 @@ func TestSimple3x3StyleOverridesAsBaseColumnRowCell(t *testing.T) {
 	assert.Equal(t, expectedTable, rendered)
 }
 
-func TestStyleFuncAppliesAfterBaseStyleAndColStylesAndBeforeRowStyle(t *testing.T) {
+func TestSimple3x3CellStyleFuncOverridesAsBaseColumnRowCell(t *testing.T) {
+	model := New([]Column{
+		NewColumn("1", "1", 6),
+		NewColumn("2", "2", 6).WithStyle(lipgloss.NewStyle().Align(lipgloss.Left)),
+		NewColumn("3", "3", 6),
+	}).WithBaseStyle(lipgloss.NewStyle().Align(lipgloss.Center))
+
+	rows := []Row{}
+
+	for rowIndex := 1; rowIndex <= 3; rowIndex++ {
+		rowData := RowData{}
+
+		for columnIndex := 1; columnIndex <= 3; columnIndex++ {
+			id := fmt.Sprintf("%d", columnIndex)
+
+			rowData[id] = fmt.Sprintf("%d,%d", columnIndex, rowIndex)
+		}
+
+		rows = append(rows, NewRow(rowData))
+	}
+
+	// Test overrides with alignment because it's easy to check output string
+	rows[0] = rows[0].WithStyle(lipgloss.NewStyle().Align(lipgloss.Left))
+	rows[0].Data["2"] = NewStyledCellWithStyleFunc("R", func(input StyledCellFuncInput) lipgloss.Style {
+		// Do some checks to make sure we're given the right information as a bonus test
+		assert.Equal(t, "2", input.Column.Key(), "Wrong column key given to style func")
+		assert.Equal(t, "R", input.Data, "Wrong data given to style func")
+
+		return lipgloss.NewStyle().Align(lipgloss.Right)
+	})
+
+	rows[2] = rows[2].WithStyle(lipgloss.NewStyle().Align(lipgloss.Right))
+
+	model = model.WithRows(rows)
+
+	const expectedTable = `┏━━━━━━┳━━━━━━┳━━━━━━┓
+┃  1   ┃2     ┃  3   ┃
+┣━━━━━━╋━━━━━━╋━━━━━━┫
+┃1,1   ┃     R┃3,1   ┃
+┃ 1,2  ┃2,2   ┃ 3,2  ┃
+┃   1,3┃   2,3┃   3,3┃
+┗━━━━━━┻━━━━━━┻━━━━━━┛`
+
+	rendered := model.View()
+
+	assert.Equal(t, expectedTable, rendered)
+}
+
+func TestRowStyleFuncAppliesAfterBaseStyleAndColStylesAndBeforeRowStyle(t *testing.T) {
 	styleFunc := func(input RowStyleFuncInput) lipgloss.Style {
 		if input.Index%2 == 0 {
 			return lipgloss.NewStyle().Align(lipgloss.Left)
@@ -849,7 +897,7 @@ func TestStyleFuncAppliesAfterBaseStyleAndColStylesAndBeforeRowStyle(t *testing.
 	assert.Equal(t, expectedTable, rendered)
 }
 
-func TestStyleFuncAppliesHighlighted(t *testing.T) {
+func TestRowStyleFuncAppliesHighlighted(t *testing.T) {
 	styleFunc := func(input RowStyleFuncInput) lipgloss.Style {
 		if input.IsHighlighted {
 			return lipgloss.NewStyle().Align(lipgloss.Center)