Allow row painting by index

table/render_init.go

@@ -223,7 +223,7 @@ func (t *Table) initForRenderRows() {
 	t.autoIndexVIndexMaxLength = len(fmt.Sprint(len(t.rowsRaw)))
 
 	// stringify all the rows to make it easy to render
-	if t.rowPainter != nil {
+	if t.rowPainter != nil || t.indexedRowPainter != nil {
 		t.rowsColors = make([]text.Colors, len(t.rowsRaw))
 	}
 	t.rows = t.initForRenderRowsStringify(t.rowsRaw, renderHint{})
@@ -277,6 +277,9 @@ func (t *Table) initForRenderSortRows() {
 		sortedRowsColors := make([]text.Colors, len(t.rows))
 		for idx := range t.rows {
 			sortedRowsColors[idx] = t.rowsColors[sortedRowIndices[idx]]
+			if t.indexedRowPainter != nil {
+				sortedRowsColors[idx] = t.indexedRowPainter(idx)
+			}
 		}
 		t.rowsColors = sortedRowsColors
 	}

table/render_test.go

@@ -879,11 +879,7 @@ func TestTable_Render_Reset(t *testing.T) {
 
 func TestTable_Render_RowPainter(t *testing.T) {
 	tw := NewWriter()
-	tw.AppendHeader(testHeader)
-	tw.AppendRows(testRows)
-	tw.AppendRow(testRowMultiLine)
-	tw.AppendFooter(testFooter)
-	tw.SetIndexColumn(1)
+
 	tw.SetRowPainter(func(row Row) text.Colors {
 		if salary, ok := row[3].(int); ok {
 			if salary > 3000 {
@@ -894,6 +890,33 @@ func TestTable_Render_RowPainter(t *testing.T) {
 		}
 		return nil
 	})
+
+	RunTestTable_Render_WithRowPainter(t, tw)
+}
+
+func TestTable_Render_IndexedRowPainter(t *testing.T) {
+	tw := NewWriter()
+
+	tw.SetIndexedRowPainter(func(idx int) text.Colors {
+		if idx > -1 {
+			if idx == 3 {
+				return text.Colors{text.BgYellow, text.FgBlack}
+			} else if idx == 0 {
+				return text.Colors{text.BgRed, text.FgBlack}
+			}
+		}
+		return nil
+	})
+
+	RunTestTable_Render_WithRowPainter(t, tw)
+}
+
+func RunTestTable_Render_WithRowPainter(t *testing.T, tw Writer) {
+	tw.AppendHeader(testHeader)
+	tw.AppendRows(testRows)
+	tw.AppendRow(testRowMultiLine)
+	tw.AppendFooter(testFooter)
+	tw.SetIndexColumn(1)
 	tw.SetStyle(StyleLight)
 	tw.SortBy([]SortBy{{Name: "Salary", Mode: AscNumeric}})
 

table/sort.go

@@ -55,6 +55,10 @@ type rowsSorter struct {
 	sortedIndices []int
 }
 
+func (t *Table) SortedIndices() []int {
+	return t.sortedRowIndices
+}
+
 // getSortedRowIndices sorts and returns the row indices in Sorted order as
 // directed by Table.sortBy which can be set using Table.SortBy(...)
 func (t *Table) getSortedRowIndices() []int {
@@ -71,6 +75,8 @@ func (t *Table) getSortedRowIndices() []int {
 		})
 	}
 
+	t.sortedRowIndices = sortedIndices
+
 	return sortedIndices
 }
 

table/table.go

@@ -26,6 +26,10 @@ func (r Row) findColumnNumber(colName string) int {
 // text.Colors{} to use on the entire row
 type RowPainter func(row Row) text.Colors
 
+// IndexedRowPainter is a custom function that takes a Row index as input and returns the
+// text.Colors{} to use on the entire row
+type IndexedRowPainter func(idx int) text.Colors
+
 // rowStr defines a single row in the Table comprised of just string objects.
 type rowStr []string
 
@@ -95,6 +99,9 @@ type Table struct {
 	// rowPainter is a custom function that given a Row, returns the colors to
 	// use on the entire row
 	rowPainter RowPainter
+	// indexedRowPainter is a custom function that given a row index, returns the colors to
+	// use on the entire row
+	indexedRowPainter IndexedRowPainter
 	// rowSeparator is a dummy row that contains the separator columns (dashes
 	// that make up the separator between header/body/footer
 	rowSeparator rowStr
@@ -103,6 +110,8 @@ type Table struct {
 	separators map[int]bool
 	// sortBy stores a map of Column
 	sortBy []SortBy
+	// sortedRowIndices contains the sorted row indices for later reference
+	sortedRowIndices []int
 	// style contains all the strings used to draw the table, and more
 	style *Style
 	// suppressEmptyColumns hides columns which have no content on all regular
@@ -317,6 +326,14 @@ func (t *Table) SetRowPainter(painter RowPainter) {
 	t.rowPainter = painter
 }
 
+// SetIndexedRowPainter sets the IndexedRowPainter function which determines the colors to use
+// on a row at index idx. Before rendering, this function is invoked on all rows and the
+// color of each row is determined. This color takes precedence over other ways
+// to set color (ColumnConfig.Color*, SetColor*()).
+func (t *Table) SetIndexedRowPainter(painter IndexedRowPainter) {
+	t.indexedRowPainter = painter
+}
+
 // SetStyle overrides the DefaultStyle with the provided one.
 func (t *Table) SetStyle(style Style) {
 	t.style = &style
@@ -461,7 +478,7 @@ func (t *Table) getColumnColors(colIdx int, hint renderHint) text.Colors {
 			return colors
 		}
 	}
-	if t.rowPainter != nil && hint.isRegularNonSeparatorRow() && !t.isIndexColumn(colIdx, hint) {
+	if (t.rowPainter != nil || t.indexedRowPainter != nil) && hint.isRegularNonSeparatorRow() && !t.isIndexColumn(colIdx, hint) {
 		if colors := t.rowsColors[hint.rowNumber-1]; colors != nil {
 			return colors
 		}

table/writer.go

@@ -13,6 +13,7 @@ type Writer interface {
 	AppendSeparator()
 	ImportGrid(grid interface{}) bool
 	Length() int
+	SortedIndices() []int
 	Pager(opts ...PagerOption) Pager
 	Render() string
 	RenderCSV() string
@@ -28,6 +29,7 @@ type Writer interface {
 	SetIndexColumn(colNum int)
 	SetOutputMirror(mirror io.Writer)
 	SetRowPainter(painter RowPainter)
+	SetIndexedRowPainter(painter IndexedRowPainter)
 	SetStyle(style Style)
 	SetTitle(format string, a ...interface{})
 	SortBy(sortBy []SortBy)