Extract table printer from CLI

example_gh_test.go

@@ -8,6 +8,8 @@ import (
 	"time"
 
 	"github.com/cli/go-gh/pkg/api"
+	"github.com/cli/go-gh/pkg/tableprinter"
+	"github.com/cli/go-gh/pkg/term"
 	graphql "github.com/cli/shurcooL-graphql"
 )
 
@@ -162,3 +164,25 @@ func ExampleCurrentRepository() {
 	}
 	fmt.Printf("%s/%s/%s\n", repo.Host(), repo.Owner(), repo.Name())
 }
+
+// Print tabular data to a terminal or in machine-readable format for scripts.
+func ExampleTablePrinter() {
+	terminal := term.FromEnv()
+	termWidth, _, _ := terminal.Size()
+	t := tableprinter.New(terminal.Out(), terminal.IsTerminalOutput(), termWidth)
+
+	red := func(s string) string {
+		return "\x1b[31m" + s + "\x1b[m"
+	}
+
+	// add a field that will render with color and will not be auto-truncated
+	t.AddField("1", tableprinter.WithColor(red), tableprinter.WithTruncate(nil))
+	t.AddField("hello")
+	t.EndRow()
+	t.AddField("2")
+	t.AddField("world")
+	t.EndRow()
+	if err := t.Render(); err != nil {
+		log.Fatal(err)
+	}
+}

go.mod

@@ -8,8 +8,11 @@ require (
 	github.com/cli/shurcooL-graphql v0.0.1
 	github.com/henvic/httpretty v0.0.6
 	github.com/kr/pretty v0.1.0 // indirect
+	github.com/mattn/go-runewidth v0.0.13
 	github.com/stretchr/testify v1.7.0
 	github.com/thlib/go-timezone-local v0.0.0-20210907160436-ef149e42d28e
+	golang.org/x/sys v0.0.0-20210831042530-f4d43177bf5e
+	golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1
 	gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect
 	gopkg.in/h2non/gock.v1 v1.1.2
 	gopkg.in/yaml.v3 v3.0.1

go.sum

@@ -15,10 +15,14 @@ github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORN
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
+github.com/mattn/go-runewidth v0.0.13 h1:lTGmDsbAYt5DmK6OnoV7EuIF1wEIFAcxld6ypU4OSgU=
+github.com/mattn/go-runewidth v0.0.13/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/nbio/st v0.0.0-20140626010706-e9e8d9816f32 h1:W6apQkHrMkS0Muv8G/TipAy/FJl/rCYT0+EuS8+Z0z4=
 github.com/nbio/st v0.0.0-20140626010706-e9e8d9816f32/go.mod h1:9wM+0iRr9ahx58uYLpLIr5fm8diHn0JbqRycJi6w0Ms=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
+github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
@@ -30,6 +34,7 @@ golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210831042530-f4d43177bf5e h1:XMgFehsDnnLGtjvjOfqWSUzt0alpTR1RSEuznObga2c=
 golang.org/x/sys v0.0.0-20210831042530-f4d43177bf5e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1 h1:v+OssWQX+hTHEmOBgwxdZxK4zHq3yOs8F9J7mk0PY8E=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

pkg/tableprinter/table.go

@@ -0,0 +1,242 @@
+// Package tableprinter facilitates rendering column-formatted data to a terminal and TSV-formatted data to
+// a script or a file. It is suitable for presenting tabular data in a human-readable format that is
+// guaranteed to fit within the given viewport, while at the same time offering the same data in a
+// machine-readable format for scripts.
+package tableprinter
+
+import (
+	"fmt"
+	"io"
+	"strings"
+
+	"github.com/mattn/go-runewidth"
+)
+
+type fieldOption func(*tableField)
+
+type TablePrinter interface {
+	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.
+func WithTruncate(fn func(int, string) string) fieldOption {
+	return func(f *tableField) {
+		f.truncateFunc = 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.
+func WithColor(fn func(string) string) fieldOption {
+	return func(f *tableField) {
+		f.colorFunc = fn
+	}
+}
+
+// New initializes a table printer with terminal mode and terminal width. When terminal mode is enabled, the
+// output will be human-readable, column-formatted to fit available width, and rendered with color support.
+// In non-terminal mode, the output is tab-separated and all truncation of values is disabled.
+func New(w io.Writer, isTTY bool, maxWidth int) TablePrinter {
+	if isTTY {
+		return &ttyTablePrinter{
+			out:      w,
+			maxWidth: maxWidth,
+		}
+	}
+	return &tsvTablePrinter{
+		out: w,
+	}
+}
+
+type tableField struct {
+	text         string
+	truncateFunc func(int, string) string
+	colorFunc    func(string) string
+}
+
+type ttyTablePrinter struct {
+	out      io.Writer
+	maxWidth int
+	rows     [][]tableField
+}
+
+func (t *ttyTablePrinter) AddField(s string, opts ...fieldOption) {
+	if t.rows == nil {
+		t.rows = make([][]tableField, 1)
+	}
+	rowI := len(t.rows) - 1
+	field := tableField{
+		text:         s,
+		truncateFunc: truncateText,
+	}
+	for _, opt := range opts {
+		opt(&field)
+	}
+	t.rows[rowI] = append(t.rows[rowI], field)
+}
+
+func (t *ttyTablePrinter) EndRow() {
+	t.rows = append(t.rows, []tableField{})
+}
+
+func (t *ttyTablePrinter) Render() error {
+	if len(t.rows) == 0 {
+		return nil
+	}
+
+	delim := "  "
+	numCols := len(t.rows[0])
+	colWidths := t.calculateColumnWidths(len(delim))
+
+	for _, row := range t.rows {
+		for col, field := range row {
+			if col > 0 {
+				_, err := fmt.Fprint(t.out, delim)
+				if err != nil {
+					return err
+				}
+			}
+			truncVal := field.text
+			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] - displayWidth(field.text); padWidth > 0 {
+					truncVal += strings.Repeat(" ", padWidth)
+				}
+			}
+			if field.colorFunc != nil {
+				truncVal = field.colorFunc(truncVal)
+			}
+			_, err := fmt.Fprint(t.out, truncVal)
+			if err != nil {
+				return err
+			}
+		}
+		if len(row) > 0 {
+			_, err := fmt.Fprint(t.out, "\n")
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (t *ttyTablePrinter) calculateColumnWidths(delimSize int) []int {
+	numCols := len(t.rows[0])
+	maxColWidths := make([]int, numCols)
+	colWidths := make([]int, numCols)
+
+	for _, row := range t.rows {
+		for col, field := range row {
+			w := displayWidth(field.text)
+			if w > maxColWidths[col] {
+				maxColWidths[col] = w
+			}
+			// if this field has disabled truncating, ensure that the column is wide enough
+			if field.truncateFunc == nil && w > colWidths[col] {
+				colWidths[col] = w
+			}
+		}
+	}
+
+	availWidth := func() int {
+		setWidths := 0
+		for col := 0; col < numCols; col++ {
+			setWidths += colWidths[col]
+		}
+		return t.maxWidth - delimSize*(numCols-1) - setWidths
+	}
+	numFixedCols := func() int {
+		fixedCols := 0
+		for col := 0; col < numCols; col++ {
+			if colWidths[col] > 0 {
+				fixedCols++
+			}
+		}
+		return fixedCols
+	}
+
+	// set the widths of short columns
+	if w := availWidth(); w > 0 {
+		if numFlexColumns := numCols - numFixedCols(); numFlexColumns > 0 {
+			perColumn := w / numFlexColumns
+			for col := 0; col < numCols; col++ {
+				if max := maxColWidths[col]; max < perColumn {
+					colWidths[col] = max
+				}
+			}
+		}
+	}
+
+	// truncate long columns to the remaining available width
+	if numFlexColumns := numCols - numFixedCols(); numFlexColumns > 0 {
+		perColumn := availWidth() / numFlexColumns
+		for col := 0; col < numCols; col++ {
+			if colWidths[col] == 0 {
+				if max := maxColWidths[col]; max < perColumn {
+					colWidths[col] = max
+				} else if perColumn > 0 {
+					colWidths[col] = perColumn
+				}
+			}
+		}
+	}
+
+	// add the remainder to truncated columns
+	if w := availWidth(); w > 0 {
+		for col := 0; col < numCols; col++ {
+			d := maxColWidths[col] - colWidths[col]
+			toAdd := w
+			if d < toAdd {
+				toAdd = d
+			}
+			colWidths[col] += toAdd
+			w -= toAdd
+			if w <= 0 {
+				break
+			}
+		}
+	}
+
+	return colWidths
+}
+
+type tsvTablePrinter struct {
+	out        io.Writer
+	currentCol int
+}
+
+func (t *tsvTablePrinter) AddField(text string, opts ...fieldOption) {
+	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
+}
+
+func truncateText(maxWidth int, s string) string {
+	if maxWidth < 5 {
+		return runewidth.Truncate(s, maxWidth, "")
+	}
+	return runewidth.Truncate(s, maxWidth, "...")
+}
+
+func displayWidth(s string) int {
+	return runewidth.StringWidth(s)
+}