From f6d33d097212cf2ea10039504c674789ecb74522 Mon Sep 17 00:00:00 2001
From: Liam Galvin <liam.galvin@aquasec.com>
Date: Tue, 19 Jul 2022 17:39:30 +0100
Subject: [PATCH] fix: Fix alignment bug with colspans (#16)

---
 table.go      | 44 ++++++++++++++++++++++++++++++++++++++++----
 table_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/table.go b/table.go
index 7335a5f..a70b635 100644
--- a/table.go
+++ b/table.go
@@ -603,25 +603,61 @@ func (t *Table) applyColSpans(formatted []iRow) []iRow {
 			target.width = targetWidth
 			formatted[job.row].cols[realTargetColIndex] = target
 		}
+	}
 
+	var lastColMaxWidth int
+	for r, row := range formatted {
+		col := row.cols[len(row.cols)-1]
+		if t.getColspan(row.header, row.footer, r, len(row.cols)-1) > 1 {
+			continue
+		}
+		width := col.width
+		if width > lastColMaxWidth {
+			lastColMaxWidth = width
+		}
+	}
+	for r, row := range formatted {
+		c := len(row.cols) - 1
+		if t.getColspan(row.header, row.footer, r, c) > 1 {
+			continue
+		}
+		width := row.cols[c].width
+		if width < lastColMaxWidth {
+			row.cols[c].width = lastColMaxWidth
+			for k, line := range row.cols[c].lines {
+				row.cols[c].lines[k] = align(line, lastColMaxWidth, row.cols[c].alignment)
+			}
+			formatted[r] = row
+		}
 	}
+
 	return formatted
 }
 
 func (t *Table) mergeContent(formatted []iRow) []iRow {
 
+	columnCount := t.calcColumnWidth(0, formatted[0])
+	lastValues := make([]string, columnCount)
+
 	// flag cols as mergeAbove where content matches and is non-empty
-	for c := 0; c < len(formatted[0].cols); c++ {
-		var previousContent string
+	for c := 0; c < columnCount; c++ {
 		var prevHeader bool
 		var allowed bool
 		for r, row := range formatted {
+			// don't merge columns with colspan > 1
+			if t.getColspan(row.header, row.footer, r, c) > 1 {
+				continue
+			}
+			if c >= len(row.cols) {
+				continue
+			}
+			relativeIndex := t.getRelativeIndex(row, c)
 			allowed = (row.header && t.autoMergeHeaders) || (!row.header && !row.footer && !prevHeader && t.autoMerge)
 			prevHeader = row.header
 			current := row.cols[c].original
-			merge := current == previousContent && strings.TrimSpace(current) != ""
+			merge := current == lastValues[relativeIndex] && strings.TrimSpace(current) != ""
 			row.cols[c].mergeAbove = merge && allowed
-			previousContent = current
+			lastValues[relativeIndex] = current
 			formatted[r] = row
 		}
 	}
diff --git a/table_test.go b/table_test.go
index 165f0fd..ea8d435 100644
--- a/table_test.go
+++ b/table_test.go
@@ -711,3 +711,30 @@ func Test_RelativeColIndexesWithSpans(t *testing.T) {
 	assert.Equal(t, 2, table.getRelativeIndex(row, 1))
 	assert.Equal(t, 5, table.getRelativeIndex(row, 2))
 }
+
+func Test_HeaderColSpanVariation(t *testing.T) {
+	builder := &strings.Builder{}
+	table := New(builder)
+	table.SetHeaders("Service", "Misconfigurations", "Last Scanned")
+	table.AddHeaders("Service", "Critical", "High", "Medium", "Low", "Unknown", "Last Scanned")
+	table.SetRowLines(false)
+	table.SetHeaderAlignment(AlignLeft, AlignCenter, AlignCenter, AlignCenter, AlignCenter, AlignCenter, AlignLeft)
+	table.SetAlignment(AlignLeft, AlignRight, AlignRight, AlignRight, AlignRight, AlignRight, AlignLeft)
+	table.SetAutoMergeHeaders(true)
+	table.SetHeaderColSpans(0, 1, 5, 1)
+	table.AddRow("ec2", "1", "2", "5", "0", "3", "2 hours ago")
+	table.AddRow("ecs", "0", "-", "-", "1", "0", "just now")
+	table.AddRow("eks", "7", "0", "0", "3", "0", "127 hours ago")
+	table.Render()
+	assertMultilineEqual(t, `
+┌─────────┬──────────────────────────────────────────────────┬───────────────┐
+│ Service │                Misconfigurations                 │ Last Scanned  │
+│         ├──────────┬──────────────┬────────┬─────┬─────────┤               │
+│         │ Critical │     High     │ Medium │ Low │ Unknown │               │
+├─────────┼──────────┼──────────────┼────────┼─────┼─────────┼───────────────┤
+│ ec2     │        1 │            2 │      5 │   0 │       3 │ 2 hours ago   │
+│ ecs     │        0 │            - │      - │   1 │       0 │ just now      │
+│ eks     │        7 │            0 │      0 │   3 │       0 │ 127 hours ago │
+└─────────┴──────────┴──────────────┴────────┴─────┴─────────┴───────────────┘
+`, "\n"+builder.String())
+}