From fce553e800176197145acff5e65cafe07c2cb650 Mon Sep 17 00:00:00 2001
From: Liam Galvin <liam.galvin@aquasec.com>
Date: Mon, 15 Aug 2022 15:01:51 +0100
Subject: [PATCH] feat: Add option to fill terminal width (#20)

---
 table.go      | 36 +++++++++++++++++++++++++++++++++---
 table_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/table.go b/table.go
index d326fe7..0867d4d 100644
--- a/table.go
+++ b/table.go
@@ -36,6 +36,7 @@ type Table struct {
 	footerColspans      map[int][]int
 	availableWidth      int
 	headerVerticalAlign Alignment
+	fillWidth           bool
 }
 
 type iRow struct {
@@ -162,6 +163,11 @@ func (t *Table) SetAutoMerge(enabled bool) {
 	t.autoMerge = enabled
 }
 
+// SetFillWidth sets whether to fill the entire available width
+func (t *Table) SetFillWidth(enabled bool) {
+	t.fillWidth = enabled
+}
+
 // SetAutoMergeHeaders sets whether to merge header cells vertically if their content is the same and non-empty
 func (t *Table) SetAutoMergeHeaders(enabled bool) {
 	t.autoMergeHeaders = enabled
@@ -232,6 +238,11 @@ func (t *Table) AddRow(cols ...string) {
 	t.data = append(t.data, cols)
 }
 
+// SetAvailableWidth sets the available width for the table (defaults to terminal width when stdout is used)
+func (t *Table) SetAvailableWidth(w int) {
+	t.availableWidth = w
+}
+
 // AddRows adds multiple rows to the table. Each argument is a row, i.e. a slice of column values.
 func (t *Table) AddRows(rows ...[]string) {
 	t.data = append(t.data, rows...)
@@ -468,21 +479,40 @@ func (t *Table) formatContent(formatted []iRow) []iRow {
 		formatted[r].height = maxLines
 	}
 
+	spares := make([]int, len(formatted))
+	for r, row := range formatted {
+		spare := t.availableWidth - 1
+		for c, col := range row.cols {
+			spare -= col.MaxWidth() + (t.getColspan(row.header, row.footer, r, c) * ((t.padding * 2) + 1))
+		}
+		if spare < 0 {
+			spare = 0
+		}
+		spares[r] = spare
+	}
+
+	var extra int
+
 	// set width of each col, and align text
 	for c := 0; c < t.calcColumnWidth(0, formatted[0]); c++ { // for each col
 
 		// find max width for column across all rows
 		maxWidth := 0
-		for _, row := range formatted {
+		for r, row := range formatted {
 
 			if c >= len(row.cols) || row.cols[c].span > 1 {
 				// ignore columns with a colspan > 1 for now, we'll apply those next
 				continue
 			}
 
-			if row.cols[c].MaxWidth() > maxWidth {
-				maxWidth = row.cols[c].MaxWidth()
+			if t.fillWidth {
+				extra = spares[r] / len(row.cols)
+			}
+			width := row.cols[c].MaxWidth() + extra
+			if width > maxWidth {
+				maxWidth = width
 			}
+
 		}
 
 		// set uniform col width, and align all content
diff --git a/table_test.go b/table_test.go
index 408796d..f1f6c6b 100644
--- a/table_test.go
+++ b/table_test.go
@@ -766,3 +766,51 @@ func Test_HeaderVerticalAlignBottom(t *testing.T) {
 └─────────┴──────────┴──────────────┴────────┴─────┴─────────┴───────────────┘
 `, "\n"+builder.String())
 }
+
+func Test_FillWidth(t *testing.T) {
+	builder := &strings.Builder{}
+	table := New(builder)
+	table.SetHeaders("A", "B", "C")
+	table.AddRow("1", "2", "3")
+	table.AddRow("4", "5", "6")
+	table.SetAvailableWidth(19)
+	table.SetFillWidth(true)
+	table.Render()
+	assertMultilineEqual(t, `
+┌─────┬─────┬─────┐
+│  A  │  B  │  C  │
+├─────┼─────┼─────┤
+│ 1   │ 2   │ 3   │
+├─────┼─────┼─────┤
+│ 4   │ 5   │ 6   │
+└─────┴─────┴─────┘
+`, "\n"+builder.String())
+}
+
+func Test_HeaderColSpanTrivyKubernetesStyleFullWithFillWidth(t *testing.T) {
+	builder := &strings.Builder{}
+	table := New(builder)
+	table.SetHeaders("Namespace", "Resource", "Vulnerabilities", "Misconfigurations")
+	table.AddHeaders("Namespace", "Resource", "Critical", "High", "Medium", "Low", "Unknown", "Critical", "High", "Medium", "Low", "Unknown")
+	table.SetHeaderColSpans(0, 1, 1, 5, 5)
+	table.SetAutoMergeHeaders(true)
+	table.SetAvailableWidth(100)
+	table.SetFillWidth(true)
+	table.AddRow("default", "Deployment/app", "2", "5", "7", "8", "0", "0", "3", "5", "19", "0")
+	table.AddRow("default", "Ingress/test", "-", "-", "-", "-", "-", "1", "0", "2", "17", "0")
+	table.AddRow("default", "Service/test", "0", "0", "0", "1", "0", "3", "0", "4", "9", "0")
+	table.Render()
+	assertMultilineEqual(t, `
+┌──────────────┬──────────────────┬──────────────────────────────────────────┬───────────────────────────────────────────┐
+│  Namespace   │     Resource     │             Vulnerabilities              │             Misconfigurations             │
+│              │                  ├──────────┬──────┬────────┬─────┬─────────┼──────────┬──────┬────────┬──────┬─────────┤
+│              │                  │ Critical │ High │ Medium │ Low │ Unknown │ Critical │ High │ Medium │ Low  │ Unknown │
+├──────────────┼──────────────────┼──────────┼──────┼────────┼─────┼─────────┼──────────┼──────┼────────┼──────┼─────────┤
+│ default      │ Deployment/app   │ 2        │ 5    │ 7      │ 8   │ 0       │ 0        │ 3    │ 5      │ 19   │ 0       │
+├──────────────┼──────────────────┼──────────┼──────┼────────┼─────┼─────────┼──────────┼──────┼────────┼──────┼─────────┤
+│ default      │ Ingress/test     │ -        │ -    │ -      │ -   │ -       │ 1        │ 0    │ 2      │ 17   │ 0       │
+├──────────────┼──────────────────┼──────────┼──────┼────────┼─────┼─────────┼──────────┼──────┼────────┼──────┼─────────┤
+│ default      │ Service/test     │ 0        │ 0    │ 0      │ 1   │ 0       │ 3        │ 0    │ 4      │ 9    │ 0       │
+└──────────────┴──────────────────┴──────────┴──────┴────────┴─────┴─────────┴──────────┴──────┴────────┴──────┴─────────┘
+`, "\n"+builder.String())
+}