table: RowPainterWithAttributes; addresses #345 (#346)

.github/workflows/ci.yml

@@ -32,6 +32,7 @@ jobs:
         run: |
           go install github.com/fzipp/gocyclo/cmd/[email protected]
           go install github.com/mattn/[email protected]
+          go install github.com/rinchsan/gosimports/cmd/[email protected]
 
       # Run all the unit-tests
       - name: Test

Makefile

@@ -6,6 +6,7 @@ all: test bench
 
 tools:
 	go install github.com/fzipp/gocyclo/cmd/[email protected]
+	go install github.com/rinchsan/gosimports/cmd/[email protected]
 
 bench:
 	go test -bench=. -benchmem
@@ -23,16 +24,18 @@ demo-table:
 	go run cmd/demo-table/demo.go
 
 fmt:
-	go fmt $(shell go list ./...)
+	go fmt ./...
+	gosimports -w .
 
 profile:
 	sh profile.sh
 
 test: fmt vet cyclo
-	go test -cover -coverprofile=.coverprofile $(shell go list ./...)
+	go test -cover -coverprofile=.coverprofile ./...
 
 test-race:
 	go run -race ./cmd/demo-progress/demo.go
 
 vet:
-	go vet $(shell go list ./...)
+	go vet ./...
+

table/pager_test.go

@@ -1,9 +1,10 @@
 package table
 
 import (
-	"github.com/stretchr/testify/assert"
 	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestPager(t *testing.T) {

table/render_init.go

@@ -223,16 +223,16 @@ 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 {
-		t.rowsColors = make([]text.Colors, len(t.rowsRaw))
-	}
 	t.rows = t.initForRenderRowsStringify(t.rowsRaw, renderHint{})
 	t.rowsFooter = t.initForRenderRowsStringify(t.rowsFooterRaw, renderHint{isFooterRow: true})
 	t.rowsHeader = t.initForRenderRowsStringify(t.rowsHeaderRaw, renderHint{isHeaderRow: true})
 
 	// sort the rows as requested
 	t.initForRenderSortRows()
 
+	// find the row colors (if any)
+	t.initForRenderRowPainterColors()
+
 	// suppress columns without any content
 	t.initForRenderSuppressColumns()
 
@@ -243,14 +243,42 @@ func (t *Table) initForRenderRows() {
 func (t *Table) initForRenderRowsStringify(rows []Row, hint renderHint) []rowStr {
 	rowsStr := make([]rowStr, len(rows))
 	for idx, row := range rows {
-		if t.rowPainter != nil && hint.isRegularRow() {
-			t.rowsColors[idx] = t.rowPainter(row)
-		}
+		hint.rowNumber = idx + 1
 		rowsStr[idx] = t.analyzeAndStringify(row, hint)
 	}
 	return rowsStr
 }
 
+func (t *Table) initForRenderRowPainterColors() {
+	if !t.hasRowPainter() {
+		return
+	}
+
+	// generate the colors
+	t.rowsColors = make([]text.Colors, len(t.rowsRaw))
+	for idx, row := range t.rowsRaw {
+		idxColors := idx
+		if len(t.sortedRowIndices) > 0 {
+			// override with the sorted row index
+			for j := 0; j < len(t.sortedRowIndices); j++ {
+				if t.sortedRowIndices[j] == idx {
+					idxColors = j
+					break
+				}
+			}
+		}
+
+		if t.rowPainter != nil {
+			t.rowsColors[idxColors] = t.rowPainter(row)
+		} else if t.rowPainterWithAttributes != nil {
+			t.rowsColors[idxColors] = t.rowPainterWithAttributes(row, RowAttributes{
+				Number:       idx + 1,
+				NumberSorted: idxColors + 1,
+			})
+		}
+	}
+}
+
 func (t *Table) initForRenderRowSeparator() {
 	t.rowSeparator = make(rowStr, t.numColumns)
 	for colIdx, maxColumnLength := range t.maxColumnLengths {
@@ -265,21 +293,12 @@ func (t *Table) initForRenderSortRows() {
 	}
 
 	// sort the rows
-	sortedRowIndices := t.getSortedRowIndices()
+	t.sortedRowIndices = t.getSortedRowIndices()
 	sortedRows := make([]rowStr, len(t.rows))
 	for idx := range t.rows {
-		sortedRows[idx] = t.rows[sortedRowIndices[idx]]
+		sortedRows[idx] = t.rows[t.sortedRowIndices[idx]]
 	}
 	t.rows = sortedRows
-
-	// sort the rowsColors
-	if len(t.rowsColors) > 0 {
-		sortedRowsColors := make([]text.Colors, len(t.rows))
-		for idx := range t.rows {
-			sortedRowsColors[idx] = t.rowsColors[sortedRowIndices[idx]]
-		}
-		t.rowsColors = sortedRowsColors
-	}
 }
 
 func (t *Table) initForRenderSuppressColumns() {

table/render_test.go

@@ -878,13 +878,51 @@ 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 {
+	runTestWithRowPainter := func(t *testing.T, rowPainter interface{}) {
+		tw := NewWriter()
+		tw.AppendHeader(testHeader)
+		tw.AppendRows(testRows)
+		tw.AppendRow(testRowMultiLine)
+		tw.AppendFooter(testFooter)
+		tw.SetIndexColumn(1)
+		tw.SetRowPainter(rowPainter)
+		tw.SetStyle(StyleLight)
+		tw.SortBy([]SortBy{{Name: "Salary", Mode: AscNumeric}})
+
+		expectedOutLines := []string{
+			"┌─────┬────────────┬───────────┬────────┬─────────────────────────────┐",
+			"│   # │ FIRST NAME │ LAST NAME │ SALARY │                             │",
+			"├─────┼────────────┼───────────┼────────┼─────────────────────────────┤",
+			"│   0 │\x1b[41;30m Winter     \x1b[0m│\x1b[41;30m Is        \x1b[0m│\x1b[41;30m      0 \x1b[0m│\x1b[41;30m Coming.                     \x1b[0m│",
+			"│     │\x1b[41;30m            \x1b[0m│\x1b[41;30m           \x1b[0m│\x1b[41;30m        \x1b[0m│\x1b[41;30m The North Remembers!        \x1b[0m│",
+			"│     │\x1b[41;30m            \x1b[0m│\x1b[41;30m           \x1b[0m│\x1b[41;30m        \x1b[0m│\x1b[41;30m This is known.              \x1b[0m│",
+			"│  20 │ Jon        │ Snow      │   2000 │ You know nothing, Jon Snow! │",
+			"│   1 │ Arya       │ Stark     │   3000 │                             │",
+			"│ 300 │\x1b[43;30m Tyrion     \x1b[0m│\x1b[43;30m Lannister \x1b[0m│\x1b[43;30m   5000 \x1b[0m│\x1b[43;30m                             \x1b[0m│",
+			"├─────┼────────────┼───────────┼────────┼─────────────────────────────┤",
+			"│     │            │ TOTAL     │  10000 │                             │",
+			"└─────┴────────────┴───────────┴────────┴─────────────────────────────┘",
+		}
+		expectedOut := strings.Join(expectedOutLines, "\n")
+		assert.Equal(t, expectedOut, tw.Render())
+
+		tw.SetStyle(StyleColoredBright)
+		tw.Style().Color.RowAlternate = tw.Style().Color.Row
+		expectedOutLines = []string{
+			"\x1b[106;30m   # \x1b[0m\x1b[106;30m FIRST NAME \x1b[0m\x1b[106;30m LAST NAME \x1b[0m\x1b[106;30m SALARY \x1b[0m\x1b[106;30m                             \x1b[0m",
+			"\x1b[106;30m   0 \x1b[0m\x1b[41;30m Winter     \x1b[0m\x1b[41;30m Is        \x1b[0m\x1b[41;30m      0 \x1b[0m\x1b[41;30m Coming.                     \x1b[0m",
+			"\x1b[106;30m     \x1b[0m\x1b[41;30m            \x1b[0m\x1b[41;30m           \x1b[0m\x1b[41;30m        \x1b[0m\x1b[41;30m The North Remembers!        \x1b[0m",
+			"\x1b[106;30m     \x1b[0m\x1b[41;30m            \x1b[0m\x1b[41;30m           \x1b[0m\x1b[41;30m        \x1b[0m\x1b[41;30m This is known.              \x1b[0m",
+			"\x1b[106;30m  20 \x1b[0m\x1b[107;30m Jon        \x1b[0m\x1b[107;30m Snow      \x1b[0m\x1b[107;30m   2000 \x1b[0m\x1b[107;30m You know nothing, Jon Snow! \x1b[0m",
+			"\x1b[106;30m   1 \x1b[0m\x1b[107;30m Arya       \x1b[0m\x1b[107;30m Stark     \x1b[0m\x1b[107;30m   3000 \x1b[0m\x1b[107;30m                             \x1b[0m",
+			"\x1b[106;30m 300 \x1b[0m\x1b[43;30m Tyrion     \x1b[0m\x1b[43;30m Lannister \x1b[0m\x1b[43;30m   5000 \x1b[0m\x1b[43;30m                             \x1b[0m",
+			"\x1b[46;30m     \x1b[0m\x1b[46;30m            \x1b[0m\x1b[46;30m TOTAL     \x1b[0m\x1b[46;30m  10000 \x1b[0m\x1b[46;30m                             \x1b[0m",
+		}
+		expectedOut = strings.Join(expectedOutLines, "\n")
+		assert.Equal(t, expectedOut, tw.Render())
+	}
+
+	rowPainter := func(row Row) text.Colors {
 		if salary, ok := row[3].(int); ok {
 			if salary > 3000 {
 				return text.Colors{text.BgYellow, text.FgBlack}
@@ -893,41 +931,25 @@ func TestTable_Render_RowPainter(t *testing.T) {
 			}
 		}
 		return nil
+	}
+	t.Run("RowPainter 1", func(t *testing.T) {
+		runTestWithRowPainter(t, rowPainter)
+	})
+	t.Run("RowPainter 2", func(t *testing.T) {
+		runTestWithRowPainter(t, RowPainter(rowPainter))
 	})
-	tw.SetStyle(StyleLight)
-	tw.SortBy([]SortBy{{Name: "Salary", Mode: AscNumeric}})
 
-	expectedOutLines := []string{
-		"┌─────┬────────────┬───────────┬────────┬─────────────────────────────┐",
-		"│   # │ FIRST NAME │ LAST NAME │ SALARY │                             │",
-		"├─────┼────────────┼───────────┼────────┼─────────────────────────────┤",
-		"│   0 │\x1b[41;30m Winter     \x1b[0m│\x1b[41;30m Is        \x1b[0m│\x1b[41;30m      0 \x1b[0m│\x1b[41;30m Coming.                     \x1b[0m│",
-		"│     │\x1b[41;30m            \x1b[0m│\x1b[41;30m           \x1b[0m│\x1b[41;30m        \x1b[0m│\x1b[41;30m The North Remembers!        \x1b[0m│",
-		"│     │\x1b[41;30m            \x1b[0m│\x1b[41;30m           \x1b[0m│\x1b[41;30m        \x1b[0m│\x1b[41;30m This is known.              \x1b[0m│",
-		"│  20 │ Jon        │ Snow      │   2000 │ You know nothing, Jon Snow! │",
-		"│   1 │ Arya       │ Stark     │   3000 │                             │",
-		"│ 300 │\x1b[43;30m Tyrion     \x1b[0m│\x1b[43;30m Lannister \x1b[0m│\x1b[43;30m   5000 \x1b[0m│\x1b[43;30m                             \x1b[0m│",
-		"├─────┼────────────┼───────────┼────────┼─────────────────────────────┤",
-		"│     │            │ TOTAL     │  10000 │                             │",
-		"└─────┴────────────┴───────────┴────────┴─────────────────────────────┘",
+	rowPainterWithAttributes := func(row Row, attr RowAttributes) text.Colors {
+		assert.NotZero(t, attr.Number)
+		assert.NotZero(t, attr.NumberSorted)
+		return rowPainter(row)
 	}
-	expectedOut := strings.Join(expectedOutLines, "\n")
-	assert.Equal(t, expectedOut, tw.Render())
-
-	tw.SetStyle(StyleColoredBright)
-	tw.Style().Color.RowAlternate = tw.Style().Color.Row
-	expectedOutLines = []string{
-		"\x1b[106;30m   # \x1b[0m\x1b[106;30m FIRST NAME \x1b[0m\x1b[106;30m LAST NAME \x1b[0m\x1b[106;30m SALARY \x1b[0m\x1b[106;30m                             \x1b[0m",
-		"\x1b[106;30m   0 \x1b[0m\x1b[41;30m Winter     \x1b[0m\x1b[41;30m Is        \x1b[0m\x1b[41;30m      0 \x1b[0m\x1b[41;30m Coming.                     \x1b[0m",
-		"\x1b[106;30m     \x1b[0m\x1b[41;30m            \x1b[0m\x1b[41;30m           \x1b[0m\x1b[41;30m        \x1b[0m\x1b[41;30m The North Remembers!        \x1b[0m",
-		"\x1b[106;30m     \x1b[0m\x1b[41;30m            \x1b[0m\x1b[41;30m           \x1b[0m\x1b[41;30m        \x1b[0m\x1b[41;30m This is known.              \x1b[0m",
-		"\x1b[106;30m  20 \x1b[0m\x1b[107;30m Jon        \x1b[0m\x1b[107;30m Snow      \x1b[0m\x1b[107;30m   2000 \x1b[0m\x1b[107;30m You know nothing, Jon Snow! \x1b[0m",
-		"\x1b[106;30m   1 \x1b[0m\x1b[107;30m Arya       \x1b[0m\x1b[107;30m Stark     \x1b[0m\x1b[107;30m   3000 \x1b[0m\x1b[107;30m                             \x1b[0m",
-		"\x1b[106;30m 300 \x1b[0m\x1b[43;30m Tyrion     \x1b[0m\x1b[43;30m Lannister \x1b[0m\x1b[43;30m   5000 \x1b[0m\x1b[43;30m                             \x1b[0m",
-		"\x1b[46;30m     \x1b[0m\x1b[46;30m            \x1b[0m\x1b[46;30m TOTAL     \x1b[0m\x1b[46;30m  10000 \x1b[0m\x1b[46;30m                             \x1b[0m",
-	}
-	expectedOut = strings.Join(expectedOutLines, "\n")
-	assert.Equal(t, expectedOut, tw.Render())
+	t.Run("RowPainterWithAttributes 1", func(t *testing.T) {
+		runTestWithRowPainter(t, rowPainterWithAttributes)
+	})
+	t.Run("RowPainterWithAttributes 2", func(t *testing.T) {
+		runTestWithRowPainter(t, RowPainterWithAttributes(rowPainterWithAttributes))
+	})
 }
 
 func TestTable_Render_Sorted(t *testing.T) {

table/row.go

@@ -0,0 +1,43 @@
+package table
+
+import (
+	"fmt"
+
+	"github.com/jedib0t/go-pretty/v6/text"
+)
+
+// Row defines a single row in the Table.
+type Row []interface{}
+
+func (r Row) findColumnNumber(colName string) int {
+	for colIdx, col := range r {
+		if fmt.Sprint(col) == colName {
+			return colIdx + 1
+		}
+	}
+	return 0
+}
+
+// RowAttributes contains properties about the Row during the render.
+type RowAttributes struct {
+	Number       int // Row Number (1-indexed) as appended
+	NumberSorted int // Row number (1-indexed) after sorting
+}
+
+// RowPainter is a custom function that takes a Row as input and returns the
+// text.Colors{} to use on the entire row
+type RowPainter func(row Row) text.Colors
+
+// RowPainterWithAttributes is the same as RowPainter but passes in additional
+// attributes from render time
+type RowPainterWithAttributes func(row Row, attr RowAttributes) text.Colors
+
+// rowStr defines a single row in the Table comprised of just string objects.
+type rowStr []string
+
+// areEqual returns true if the contents of the 2 given columns are the same
+func (row rowStr) areEqual(colIdx1 int, colIdx2 int) bool {
+	return colIdx1 >= 0 && colIdx1 < len(row) &&
+		colIdx2 >= 0 && colIdx2 < len(row) &&
+		row[colIdx1] == row[colIdx2]
+}