Add header support to tableprinter (#139)

* Add header support to tableprinter

* Define AddHeaders instead of WithOptions pattern

Resolves PR feedback

* Add WithPadding

Resolves PR feedback

* Add text.PadRight function

---------

Co-authored-by: Sam Coe <samcoe@users.noreply.github.com>

pkg/tableprinter/table.go

@@ -7,30 +7,42 @@ package tableprinter
 import (
 	"fmt"
 	"io"
-	"strings"
 
 	"github.com/cli/go-gh/v2/pkg/text"
 )
 
 type fieldOption func(*tableField)
 
 type TablePrinter interface {
+	AddHeader([]string, ...fieldOption)
 	AddField(string, ...fieldOption)
 	EndRow()
 	Render() error
 }
 
 // WithTruncate overrides the truncation function for the field. The function should transform a string
 // argument into a string that fits within the given display width. The default behavior is to truncate the
-// value by adding "..." in the end. Pass nil to disable truncation for this value.
+// value by adding "..." in the end. The truncation function will be called before padding and coloring.
+// Pass nil to disable truncation for this value.
 func WithTruncate(fn func(int, string) string) fieldOption {
 	return func(f *tableField) {
 		f.truncateFunc = fn
 	}
 }
 
+// WithPadding overrides the padding function for the field. The function should transform a string argument
+// into a string that is padded to fit within the given display width. The default behavior is to pad fields
+// with spaces except for the last field. The padding function will be called after truncation and before coloring.
+// Pass nil to disable padding for this value.
+func WithPadding(fn func(int, string) string) fieldOption {
+	return func(f *tableField) {
+		f.paddingFunc = fn
+	}
+}
+
 // WithColor sets the color function for the field. The function should transform a string value by wrapping
 // it in ANSI escape codes. The color function will not be used if the table was initialized in non-terminal mode.
+// The color function will be called before truncation and padding.
 func WithColor(fn func(string) string) fieldOption {
 	return func(f *tableField) {
 		f.colorFunc = fn
@@ -47,6 +59,7 @@ func New(w io.Writer, isTTY bool, maxWidth int) TablePrinter {
 			maxWidth: maxWidth,
 		}
 	}
+
 	return &tsvTablePrinter{
 		out: w,
 	}
@@ -55,13 +68,27 @@ func New(w io.Writer, isTTY bool, maxWidth int) TablePrinter {
 type tableField struct {
 	text         string
 	truncateFunc func(int, string) string
+	paddingFunc  func(int, string) string
 	colorFunc    func(string) string
 }
 
 type ttyTablePrinter struct {
-	out      io.Writer
-	maxWidth int
-	rows     [][]tableField
+	out        io.Writer
+	maxWidth   int
+	hasHeaders bool
+	rows       [][]tableField
+}
+
+func (t *ttyTablePrinter) AddHeader(columns []string, opts ...fieldOption) {
+	if t.hasHeaders {
+		return
+	}
+
+	t.hasHeaders = true
+	for _, column := range columns {
+		t.AddField(column, opts...)
+	}
+	t.EndRow()
 }
 
 func (t *ttyTablePrinter) AddField(s string, opts ...fieldOption) {
@@ -104,11 +131,10 @@ func (t *ttyTablePrinter) Render() error {
 			if field.truncateFunc != nil {
 				truncVal = field.truncateFunc(colWidths[col], field.text)
 			}
-			if col < numCols-1 {
-				// pad value with spaces on the right
-				if padWidth := colWidths[col] - text.DisplayWidth(field.text); padWidth > 0 {
-					truncVal += strings.Repeat(" ", padWidth)
-				}
+			if field.paddingFunc != nil {
+				truncVal = field.paddingFunc(colWidths[col], truncVal)
+			} else if col < numCols-1 {
+				truncVal = text.PadRight(colWidths[col], truncVal)
 			}
 			if field.colorFunc != nil {
 				truncVal = field.colorFunc(truncVal)
@@ -213,7 +239,9 @@ type tsvTablePrinter struct {
 	currentCol int
 }
 
-func (t *tsvTablePrinter) AddField(text string, opts ...fieldOption) {
+func (t *tsvTablePrinter) AddHeader(_ []string, _ ...fieldOption) {}
+
+func (t *tsvTablePrinter) AddField(text string, _ ...fieldOption) {
 	if t.currentCol > 0 {
 		fmt.Fprint(t.out, "\t")
 	}

pkg/tableprinter/table_test.go

@@ -2,9 +2,13 @@ package tableprinter
 
 import (
 	"bytes"
+	"fmt"
 	"log"
 	"os"
+	"strings"
 	"testing"
+
+	"github.com/MakeNowJust/heredoc"
 )
 
 func ExampleTablePrinter() {
@@ -74,6 +78,64 @@ func Test_ttyTablePrinter_WithTruncate(t *testing.T) {
 	}
 }
 
+func Test_ttyTablePrinter_AddHeader(t *testing.T) {
+	buf := bytes.Buffer{}
+	tp := New(&buf, true, 80)
+
+	tp.AddHeader([]string{"ONE", "TWO", "THREE"}, WithColor(func(s string) string {
+		return fmt.Sprintf("\x1b[4m%s\x1b[m", s)
+	}))
+	// Subsequent calls to AddHeader are ignored.
+	tp.AddHeader([]string{"SHOULD", "NOT", "EXIST"})
+
+	tp.AddField("hello")
+	tp.AddField("beautiful")
+	tp.AddField("people")
+	tp.EndRow()
+
+	err := tp.Render()
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	expected := heredoc.Docf(`
+		%[1]s[4mONE  %[1]s[m  %[1]s[4mTWO      %[1]s[m  %[1]s[4mTHREE%[1]s[m
+		hello  beautiful  people
+	`, "\x1b")
+	if buf.String() != expected {
+		t.Errorf("expected: %q, got: %q", expected, buf.String())
+	}
+}
+
+func Test_ttyTablePrinter_WithPadding(t *testing.T) {
+	buf := bytes.Buffer{}
+	tp := New(&buf, true, 80)
+
+	// Center the headers.
+	tp.AddHeader([]string{"A", "B", "C"}, WithPadding(func(width int, s string) string {
+		left := (width - len(s)) / 2
+		return strings.Repeat(" ", left) + s + strings.Repeat(" ", width-left-len(s))
+	}))
+
+	tp.AddField("hello")
+	tp.AddField("beautiful")
+	tp.AddField("people")
+	tp.EndRow()
+
+	err := tp.Render()
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	expected := heredoc.Doc(`
+		  A        B        C   
+		hello  beautiful  people
+	`)
+	if buf.String() != expected {
+		t.Errorf("expected: %q, got: %q", expected, buf.String())
+	}
+}
+
 func Test_tsvTablePrinter(t *testing.T) {
 	buf := bytes.Buffer{}
 	tp := New(&buf, false, 0)
@@ -95,3 +157,30 @@ func Test_tsvTablePrinter(t *testing.T) {
 		t.Errorf("expected: %q, got: %q", expected, buf.String())
 	}
 }
+
+func Test_tsvTablePrinter_AddHeader(t *testing.T) {
+	buf := bytes.Buffer{}
+	tp := New(&buf, false, 0)
+
+	// Headers are not output in TSV output.
+	tp.AddHeader([]string{"ONE", "TWO", "THREE"})
+
+	tp.AddField("hello")
+	tp.AddField("beautiful")
+	tp.AddField("people")
+	tp.EndRow()
+	tp.AddField("1")
+	tp.AddField("2")
+	tp.AddField("3")
+	tp.EndRow()
+
+	err := tp.Render()
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	expected := "hello\tbeautiful\tpeople\n1\t2\t3\n"
+	if buf.String() != expected {
+		t.Errorf("expected: %q, got: %q", expected, buf.String())
+	}
+}

pkg/text/text.go

@@ -53,6 +53,15 @@ func Truncate(maxWidth int, s string) string {
 	return r
 }
 
+// PadRight returns a copy of the string s that has been padded on the right with whitespace to fit
+// the maximum display width.
+func PadRight(maxWidth int, s string) string {
+	if padWidth := maxWidth - DisplayWidth(s); padWidth > 0 {
+		s += strings.Repeat(" ", padWidth)
+	}
+	return s
+}
+
 // Pluralize returns a concatenated string with num and the plural form of thing if necessary.
 func Pluralize(num int, thing string) string {
 	if num == 1 {

pkg/text/text_test.go

@@ -172,6 +172,129 @@ func TestTruncate(t *testing.T) {
 	}
 }
 
+func TestPadRight(t *testing.T) {
+	type args struct {
+		max int
+		s   string
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			name: "empty",
+			args: args{
+				s:   "",
+				max: 5,
+			},
+			want: "     ",
+		},
+		{
+			name: "short",
+			args: args{
+				s:   "hello",
+				max: 7,
+			},
+			want: "hello  ",
+		},
+		{
+			name: "long",
+			args: args{
+				s:   "hello world",
+				max: 5,
+			},
+			want: "hello world",
+		},
+		{
+			name: "exact",
+			args: args{
+				s:   "hello world",
+				max: 11,
+			},
+			want: "hello world",
+		},
+		{
+			name: "Japanese",
+			args: args{
+				s:   "テストテスト",
+				max: 13,
+			},
+			want: "テストテスト ",
+		},
+		{
+			name: "Japanese filled",
+			args: args{
+				s:   "aテスト",
+				max: 9,
+			},
+			want: "aテスト  ",
+		},
+		{
+			name: "Chinese",
+			args: args{
+				s:   "幫新舉報違章工廠新增編號",
+				max: 26,
+			},
+			want: "幫新舉報違章工廠新增編號  ",
+		},
+		{
+			name: "Chinese filled",
+			args: args{
+				s:   "a幫新舉報違章工廠新增編號",
+				max: 26,
+			},
+			want: "a幫新舉報違章工廠新增編號 ",
+		},
+		{
+			name: "Korean",
+			args: args{
+				s:   "프로젝트 내의",
+				max: 15,
+			},
+			want: "프로젝트 내의  ",
+		},
+		{
+			name: "Korean filled",
+			args: args{
+				s:   "a프로젝트 내의",
+				max: 15,
+			},
+			want: "a프로젝트 내의 ",
+		},
+		{
+			name: "Emoji",
+			args: args{
+				s:   "💡💡💡💡",
+				max: 10,
+			},
+			want: "💡💡💡💡  ",
+		},
+		{
+			name: "Accented characters",
+			args: args{
+				s:   "é́́é́́é́́é́́é́́",
+				max: 7,
+			},
+			want: "é́́é́́é́́é́́é́́  ",
+		},
+		{
+			name: "Red accented characters",
+			args: args{
+				s:   "\x1b[0;31mé́́é́́é́́é́́é́́\x1b[0m",
+				max: 7,
+			},
+			want: "\x1b[0;31mé́́é́́é́́é́́é́́\x1b[0m  ",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := PadRight(tt.args.max, tt.args.s)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}
+
 func TestDisplayWidth(t *testing.T) {
 	tests := []struct {
 		name string