Redesign TablePrinter to avoid SetContentWidth / FitColumns steps

The API is now:
- AddField;
- EndRow;
- Render.

command/pr.go

@@ -167,33 +167,37 @@ func prList(cmd *cobra.Command, args []string) error {
 	}
 
 	table := utils.NewTablePrinter(cmd.OutOrStdout())
-	for _, pr := range prs {
-		table.SetContentWidth(0, len(strconv.Itoa(pr.Number))+1)
-		table.SetContentWidth(1, len(pr.Title))
-		table.SetContentWidth(2, len(pr.HeadLabel()))
-	}
-
-	table.FitColumns()
-	table.SetColorFunc(2, utils.Cyan)
-
 	for _, pr := range prs {
 		prNum := strconv.Itoa(pr.Number)
-		if table.IsTTY {
+		if table.IsTTY() {
 			prNum = "#" + prNum
 		}
-		switch pr.State {
-		case "OPEN":
-			table.SetColorFunc(0, utils.Green)
-		case "CLOSED":
-			table.SetColorFunc(0, utils.Red)
-		case "MERGED":
-			table.SetColorFunc(0, utils.Magenta)
-		}
-		table.WriteRow(prNum, pr.Title, pr.HeadLabel())
+		table.AddField(prNum, nil, colorFuncForState(pr.State))
+		table.AddField(pr.Title, nil, nil)
+		table.AddField(pr.HeadLabel(), nil, utils.Cyan)
+		table.EndRow()
+	}
+	err = table.Render()
+	if err != nil {
+		return err
 	}
+
 	return nil
 }
 
+func colorFuncForState(state string) func(string) string {
+	switch state {
+	case "OPEN":
+		return utils.Green
+	case "CLOSED":
+		return utils.Red
+	case "MERGED":
+		return utils.Magenta
+	default:
+		return nil
+	}
+}
+
 func prView(cmd *cobra.Command, args []string) error {
 	ctx := contextForCommand(cmd)
 	baseRepo, err := ctx.BaseRepo()

utils/table_printer.go

@@ -8,101 +8,157 @@ import (
 	"golang.org/x/crypto/ssh/terminal"
 )
 
-func NewTablePrinter(w io.Writer) *TTYTablePrinter {
-	tty := false
-	ttyWidth := 80
+type TablePrinter interface {
+	IsTTY() bool
+	AddField(string, func(int, string) string, func(string) string)
+	EndRow()
+	Render() error
+}
+
+func NewTablePrinter(w io.Writer) TablePrinter {
 	if outFile, isFile := w.(*os.File); isFile {
 		fd := int(outFile.Fd())
-		tty = terminal.IsTerminal(fd)
-		if w, _, err := terminal.GetSize(fd); err == nil {
-			ttyWidth = w
+		if terminal.IsTerminal(fd) {
+			ttyWidth := 80
+			if w, _, err := terminal.GetSize(fd); err == nil {
+				ttyWidth = w
+			}
+			return &ttyTablePrinter{
+				out:      w,
+				maxWidth: ttyWidth,
+			}
 		}
 	}
-	return &TTYTablePrinter{
-		out:       w,
-		IsTTY:     tty,
-		maxWidth:  ttyWidth,
-		colWidths: make(map[int]int),
-		colFuncs:  make(map[int]func(string) string),
+	return &tsvTablePrinter{
+		out: w,
 	}
 }
 
-type TTYTablePrinter struct {
-	out       io.Writer
-	IsTTY     bool
-	maxWidth  int
-	colWidths map[int]int
-	colFuncs  map[int]func(string) string
+type tableField struct {
+	Text         string
+	TruncateFunc func(int, string) string
+	ColorFunc    func(string) string
 }
 
-func (t *TTYTablePrinter) SetContentWidth(col, width int) {
-	if width > t.colWidths[col] {
-		t.colWidths[col] = width
+type ttyTablePrinter struct {
+	out      io.Writer
+	maxWidth int
+	rows     [][]tableField
+}
+
+func (t ttyTablePrinter) IsTTY() bool {
+	return true
+}
+
+func (t *ttyTablePrinter) AddField(text string, truncateFunc func(int, string) string, colorFunc func(string) string) {
+	if truncateFunc == nil {
+		truncateFunc = truncate
+	}
+	if t.rows == nil {
+		t.rows = [][]tableField{[]tableField{}}
 	}
+	rowI := len(t.rows) - 1
+	field := tableField{
+		Text:         text,
+		TruncateFunc: truncateFunc,
+		ColorFunc:    colorFunc,
+	}
+	t.rows[rowI] = append(t.rows[rowI], field)
 }
 
-func (t *TTYTablePrinter) SetColorFunc(col int, colorize func(string) string) {
-	t.colFuncs[col] = colorize
+func (t *ttyTablePrinter) EndRow() {
+	t.rows = append(t.rows, []tableField{})
 }
 
-// FitColumns caps all but first column to fit available terminal width.
-func (t *TTYTablePrinter) FitColumns() {
-	numCols := len(t.colWidths)
-	delimWidth := 2
-	availWidth := t.maxWidth - t.colWidths[0] - ((numCols - 1) * delimWidth)
+func (t *ttyTablePrinter) Render() error {
+	if len(t.rows) == 0 {
+		return nil
+	}
+
+	numCols := len(t.rows[0])
+	colWidths := make([]int, numCols)
+	// measure maximum content width per column
+	for _, row := range t.rows {
+		for col, field := range row {
+			textLen := len(field.Text)
+			if textLen > colWidths[col] {
+				colWidths[col] = textLen
+			}
+		}
+	}
+
+	delim := "  "
+	availWidth := t.maxWidth - colWidths[0] - ((numCols - 1) * len(delim))
 	// add extra space from columns that are already narrower than threshold
-	for col := 1; col < len(t.colWidths); col++ {
+	for col := 1; col < numCols; col++ {
 		availColWidth := availWidth / (numCols - 1)
-		if extra := availColWidth - t.colWidths[col]; extra > 0 {
+		if extra := availColWidth - colWidths[col]; extra > 0 {
 			availWidth += extra
 		}
 	}
+	// cap all but first column to fit available terminal width
 	// TODO: support weighted instead of even redistribution
-	for col := 1; col < len(t.colWidths); col++ {
+	for col := 1; col < numCols; col++ {
 		availColWidth := availWidth / (numCols - 1)
-		if t.colWidths[col] > availColWidth {
-			t.colWidths[col] = availColWidth
+		if colWidths[col] > availColWidth {
+			colWidths[col] = availColWidth
 		}
 	}
-}
-
-func (t *TTYTablePrinter) WriteRow(fields ...string) error {
-	lastCol := len(fields) - 1
-	delim := "\t"
-	if t.IsTTY {
-		delim = "  "
-	}
 
-	for col, val := range fields {
-		if col > 0 {
-			_, err := fmt.Fprint(t.out, delim)
-			if err != nil {
-				return err
+	for _, row := range t.rows {
+		for col, field := range row {
+			if col > 0 {
+				_, err := fmt.Fprint(t.out, delim)
+				if err != nil {
+					return err
+				}
 			}
-		}
-		if t.IsTTY {
-			truncVal := truncate(t.colWidths[col], val)
-			if col != lastCol {
-				truncVal = fmt.Sprintf("%-*s", t.colWidths[col], truncVal)
+			truncVal := field.TruncateFunc(colWidths[col], field.Text)
+			if col < numCols-1 {
+				// pad value with spaces on the right
+				truncVal = fmt.Sprintf("%-*s", colWidths[col], truncVal)
 			}
-			if t.colFuncs[col] != nil {
-				truncVal = t.colFuncs[col](truncVal)
+			if field.ColorFunc != nil {
+				truncVal = field.ColorFunc(truncVal)
 			}
 			_, err := fmt.Fprint(t.out, truncVal)
 			if err != nil {
 				return err
 			}
-		} else {
-			_, err := fmt.Fprint(t.out, val)
+		}
+		if len(row) > 0 {
+			_, err := fmt.Fprint(t.out, "\n")
 			if err != nil {
 				return err
 			}
 		}
 	}
-	_, err := fmt.Fprint(t.out, "\n")
-	if err != nil {
-		return err
+	return nil
+}
+
+type tsvTablePrinter struct {
+	out        io.Writer
+	currentCol int
+}
+
+func (t tsvTablePrinter) IsTTY() bool {
+	return false
+}
+
+func (t *tsvTablePrinter) AddField(text string, _ func(int, string) string, _ func(string) string) {
+	if t.currentCol > 0 {
+		fmt.Fprint(t.out, "\t")
 	}
+	fmt.Fprint(t.out, text)
+	t.currentCol++
+}
+
+func (t *tsvTablePrinter) EndRow() {
+	fmt.Fprint(t.out, "\n")
+	t.currentCol = 0
+}
+
+func (t *tsvTablePrinter) Render() error {
 	return nil
 }