From 3372b97dd4c05c70a8bd66f50372a002650dd89d Mon Sep 17 00:00:00 2001
From: Brandon Fulljames <bfullj@gmail.com>
Date: Thu, 14 Apr 2022 21:00:52 +0900
Subject: [PATCH] Add WithColumns option (#62)

---
 examples/updates/main.go | 39 ++++++++++++++++++++++++++++++---------
 table/options.go         | 12 ++++++++++++
 table/view_test.go       | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 9 deletions(-)

diff --git a/examples/updates/main.go b/examples/updates/main.go
index 1f0eea1..207a9ac 100644
--- a/examples/updates/main.go
+++ b/examples/updates/main.go
@@ -32,14 +32,8 @@ type Model struct {
 }
 
 func NewModel() Model {
-	columns := []table.Column{
-		table.NewColumn(columnKeyID, "ID", 10),
-		table.NewColumn(columnKeyScore, "Score", 8),
-		table.NewColumn(columnKeyStatus, "Status", 10),
-	}
-
 	return Model{
-		table:       table.New(columns),
+		table:       table.New(generateColumns(0)),
 		updateDelay: time.Second,
 	}
 }
@@ -58,6 +52,25 @@ func refreshDataCmd() tea.Msg {
 	}
 }
 
+// Generate columns based on how many are critical to show some summary
+func generateColumns(numCritical int) []table.Column {
+	// Show how many critical there are
+	statusStr := fmt.Sprintf("Score (%d)", numCritical)
+	statusCol := table.NewColumn(columnKeyStatus, statusStr, 10)
+
+	if numCritical > 3 {
+		// This normally applies the critical style to everything in the column,
+		// but in this case we apply a row style which overrides it anyway.
+		statusCol = statusCol.WithStyle(styleCritical)
+	}
+
+	return []table.Column{
+		table.NewColumn(columnKeyID, "ID", 10),
+		table.NewColumn(columnKeyScore, "Score", 8),
+		statusCol,
+	}
+}
+
 func (m Model) Init() tea.Cmd {
 	return refreshDataCmd
 }
@@ -91,8 +104,16 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	case []*SomeData:
 		m.data = msg
 
-		// Reapply the new data
-		m.table = m.table.WithRows(generateRowsFromData(m.data))
+		numCritical := 0
+
+		for _, d := range msg {
+			if d.Status == "Critical" {
+				numCritical++
+			}
+		}
+
+		// Reapply the new data and the new columns based on critical count
+		m.table = m.table.WithRows(generateRowsFromData(m.data)).WithColumns(generateColumns(numCritical))
 
 		// This can be from any source, but for demo purposes let's party!
 		delay := m.updateDelay
diff --git a/table/options.go b/table/options.go
index 439c40f..d2d596a 100644
--- a/table/options.go
+++ b/table/options.go
@@ -230,3 +230,15 @@ func (m Model) WithCurrentPage(currentPage int) Model {
 
 	return m
 }
+
+// WithColumns sets the visible columns for the table, so that columns can be
+// added/removed/resized or headers rewritten.
+func (m Model) WithColumns(columns []Column) Model {
+	// Deep copy to avoid edits
+	m.columns = make([]Column, len(columns))
+	copy(m.columns, columns)
+
+	m.recalculateWidth()
+
+	return m
+}
diff --git a/table/view_test.go b/table/view_test.go
index ce99a0b..9868f21 100644
--- a/table/view_test.go
+++ b/table/view_test.go
@@ -562,3 +562,40 @@ func TestSimpleFlex3x3AtAllTargetWidths(t *testing.T) {
 		}
 	}
 }
+
+func TestViewResizesWhenColumnsChange(t *testing.T) {
+	model := New([]Column{
+		NewColumn("id", "ID", 4),
+	}).WithRows([]Row{
+		NewRow(RowData{"id": "1", "score": 3}),
+		NewRow(RowData{"id": "2", "score": 4}),
+	})
+
+	const expectedTableOriginal = `┏━━━━┓
+┃  ID┃
+┣━━━━┫
+┃   1┃
+┃   2┃
+┗━━━━┛`
+
+	// Lowercased, resized, and new column added
+	const expectedTableUpdated = `┏━━━━━┳━━━━━━┓
+┃   id┃ Score┃
+┣━━━━━╋━━━━━━┫
+┃    1┃     3┃
+┃    2┃     4┃
+┗━━━━━┻━━━━━━┛`
+
+	rendered := model.View()
+
+	assert.Equal(t, expectedTableOriginal, rendered)
+
+	model = model.WithColumns([]Column{
+		NewColumn("id", "id", 5),
+		NewColumn("score", "Score", 6),
+	})
+
+	rendered = model.View()
+
+	assert.Equal(t, expectedTableUpdated, rendered)
+}