chore: make output formatting generic

coder · deansheather · Feb 8, 2023 · Feb 7, 2023 · Feb 7, 2023 · Feb 7, 2023
commit 457cd51898667bb049156ec49fa84e3740adcc8c
diff --git a/cli/cliui/output.go b/cli/cliui/output.go
@@ -0,0 +1,153 @@
+package cliui
+
+import (
+	"context"
+	"encoding/json"
+	"reflect"
+	"strings"
+
+	"github.com/spf13/cobra"
+	"golang.org/x/xerrors"
+)
+
+type OutputFormat interface {
+	ID() string
+	AttachFlags(cmd *cobra.Command)
+	Format(ctx context.Context, data any) (string, error)
+}
+
+type OutputFormatter struct {
+	formats  []OutputFormat
+	formatID string
+}
+
+// NewOutputFormatter creates a new OutputFormatter with the given formats. The
+// first format is the default format. At least two formats must be provided.
+func NewOutputFormatter(formats ...OutputFormat) *OutputFormatter {
+	if len(formats) < 2 {
+		panic("at least two output formats must be provided")
+	}
+
+	formatIDs := make(map[string]struct{}, len(formats))
+	for _, format := range formats {
+		if _, ok := formatIDs[format.ID()]; ok {
+			panic("duplicate format ID: " + format.ID())
+		}
+		formatIDs[format.ID()] = struct{}{}
+	}
+
+	return &OutputFormatter{
+		formats:  formats,
+		formatID: formats[0].ID(),
+	}
+}
+
+// AttachFlags attaches the --output flag to the given command, and any
+// additional flags required by the output formatters.
+func (f *OutputFormatter) AttachFlags(cmd *cobra.Command) {
+	for _, format := range f.formats {
+		format.AttachFlags(cmd)
+	}
+
+	formatNames := make([]string, 0, len(f.formats))
+	for _, format := range f.formats {
+		formatNames = append(formatNames, format.ID())
+	}
+
+	cmd.Flags().StringVarP(&f.formatID, "output", "o", f.formats[0].ID(), "Output format. Available formats: "+strings.Join(formatNames, ", "))
+}
+
+// Format formats the given data using the format specified by the --output
+// flag. If the flag is not set, the default format is used.
+func (f *OutputFormatter) Format(ctx context.Context, data any) (string, error) {
+	for _, format := range f.formats {
+		if format.ID() == f.formatID {
+			return format.Format(ctx, data)
+		}
+	}
+
+	return "", xerrors.Errorf("unknown output format %q", f.formatID)
+}
+
+type tableFormat struct {
+	defaultColumns []string
+	allColumns     []string
+	sort           string
+
+	columns []string
+}
+
+var _ OutputFormat = &tableFormat{}
+
+// TableFormat creates a table formatter for the given output type. The output
+// type should be specified as an empty slice of the desired type.
+//
+// E.g.: TableFormat([]MyType{}, []string{"foo", "bar"})
+//
+// defaultColumns is optional and specifies the default columns to display. If
+// not specified, all columns are displayed by default.
+func TableFormat(out any, defaultColumns []string) OutputFormat {
+	v := reflect.Indirect(reflect.ValueOf(out))
+	if v.Kind() != reflect.Slice {
+		panic("DisplayTable called with a non-slice type")
+	}
+
+	// Get the list of table column headers.
+	headers, defaultSort, err := typeToTableHeaders(v.Type().Elem())
+	if err != nil {
+		panic("parse table headers: " + err.Error())
+	}
+
+	tf := &tableFormat{
+		defaultColumns: headers,
+		allColumns:     headers,
+		sort:           defaultSort,
+	}
+	if len(defaultColumns) > 0 {
+		tf.defaultColumns = defaultColumns
+	}
+
+	return tf
+}
+
+// ID implements OutputFormat.
+func (*tableFormat) ID() string {
+	return "table"
+}
+
+// AttachFlags implements OutputFormat.
+func (f *tableFormat) AttachFlags(cmd *cobra.Command) {
+	cmd.Flags().StringSliceVarP(&f.columns, "column", "c", f.defaultColumns, "Columns to display in table output. Available columns: "+strings.Join(f.allColumns, ", "))
+}
+
+// Format implements OutputFormat.
+func (f *tableFormat) Format(_ context.Context, data any) (string, error) {
+	return DisplayTable(data, f.sort, f.columns)
+}
+
+type jsonFormat struct{}
+
+var _ OutputFormat = jsonFormat{}
+
+// JSONFormat creates a JSON formatter.
+func JSONFormat() OutputFormat {
+	return jsonFormat{}
+}
+
+// ID implements OutputFormat.
+func (jsonFormat) ID() string {
+	return "json"
+}
+
+// AttachFlags implements OutputFormat.
+func (jsonFormat) AttachFlags(_ *cobra.Command) {}
+
+// Format implements OutputFormat.
+func (jsonFormat) Format(_ context.Context, data any) (string, error) {
+	outBytes, err := json.MarshalIndent(data, "", "  ")
+	if err != nil {
+		return "", xerrors.Errorf("marshal output to JSON: %w", err)
+	}
+
+	return string(outBytes), nil
+}
diff --git a/cli/cliui/table.go b/cli/cliui/table.go
@@ -22,10 +22,10 @@ func Table() table.Writer {
 	return tableWriter
 }
 
-// FilterTableColumns returns configurations to hide columns
+// filterTableColumns returns configurations to hide columns
 // that are not provided in the array. If the array is empty,
 // no filtering will occur!
-func FilterTableColumns(header table.Row, columns []string) []table.ColumnConfig {
+func filterTableColumns(header table.Row, columns []string) []table.ColumnConfig {
 	if len(columns) == 0 {
 		return nil
 	}
@@ -51,6 +51,9 @@ func FilterTableColumns(header table.Row, columns []string) []table.ColumnConfig
 // of structs. At least one field in the struct must have a `table:""` tag
 // containing the name of the column in the outputted table.
 //
+// If `sort` is not specified, the field with the `table:"$NAME,default_sort"`
+// tag will be used to sort. An error will be returned if no field has this tag.
+//
 // Nested structs are processed if the field has the `table:"$NAME,recursive"`
 // tag and their fields will be named as `$PARENT_NAME $NAME`. If the tag is
 // malformed or a field is marked as recursive but does not contain a struct or
@@ -67,13 +70,16 @@ func DisplayTable(out any, sort string, filterColumns []string) (string, error)
 	}
 
 	// Get the list of table column headers.
-	headersRaw, err := typeToTableHeaders(v.Type().Elem())
+	headersRaw, defaultSort, err := typeToTableHeaders(v.Type().Elem())
 	if err != nil {
 		return "", xerrors.Errorf("get table headers recursively for type %q: %w", v.Type().Elem().String(), err)
 	}
 	if len(headersRaw) == 0 {
 		return "", xerrors.New(`no table headers found on the input type, make sure there is at least one "table" struct tag`)
 	}
+	if sort == "" {
+		sort = defaultSort
+	}
 	headers := make(table.Row, len(headersRaw))
 	for i, header := range headersRaw {
 		headers[i] = header
@@ -128,7 +134,7 @@ func DisplayTable(out any, sort string, filterColumns []string) (string, error)
 	// Setup the table formatter.
 	tw := Table()
 	tw.AppendHeader(headers)
-	tw.SetColumnConfigs(FilterTableColumns(headers, filterColumns))
+	tw.SetColumnConfigs(filterTableColumns(headers, filterColumns))
 	if sort != "" {
 		tw.SortBy([]table.SortBy{{
 			Name: sort,
@@ -182,29 +188,32 @@ func DisplayTable(out any, sort string, filterColumns []string) (string, error)
 // returned. If the table tag is malformed, an error is returned.
 //
 // The returned name is transformed from "snake_case" to "normal text".
-func parseTableStructTag(field reflect.StructField) (name string, recurse bool, err error) {
+func parseTableStructTag(field reflect.StructField) (name string, defaultSort, recursive bool, err error) {
 	tags, err := structtag.Parse(string(field.Tag))
 	if err != nil {
-		return "", false, xerrors.Errorf("parse struct field tag %q: %w", string(field.Tag), err)
+		return "", false, false, xerrors.Errorf("parse struct field tag %q: %w", string(field.Tag), err)
 	}
 
 	tag, err := tags.Get("table")
 	if err != nil || tag.Name == "-" {
 		// tags.Get only returns an error if the tag is not found.
-		return "", false, nil
+		return "", false, false, nil
 	}
 
-	recursive := false
+	defaultSortOpt := false
+	recursiveOpt := false
 	for _, opt := range tag.Options {
-		if opt == "recursive" {
-			recursive = true
-			continue
+		switch opt {
+		case "default_sort":
+			defaultSortOpt = true
+		case "recursive":
+			recursiveOpt = true
+		default:
+			return "", false, false, xerrors.Errorf("unknown option %q in struct field tag", opt)
 		}
-
-		return "", false, xerrors.Errorf("unknown option %q in struct field tag", opt)
 	}
 
-	return strings.ReplaceAll(tag.Name, "_", " "), recursive, nil
+	return strings.ReplaceAll(tag.Name, "_", " "), defaultSortOpt, recursiveOpt, nil
 }
 
 func isStructOrStructPointer(t reflect.Type) bool {
@@ -214,34 +223,41 @@ func isStructOrStructPointer(t reflect.Type) bool {
 // typeToTableHeaders converts a type to a slice of column names. If the given
 // type is invalid (not a struct or a pointer to a struct, has invalid table
 // tags, etc.), an error is returned.
-func typeToTableHeaders(t reflect.Type) ([]string, error) {
+func typeToTableHeaders(t reflect.Type) ([]string, string, error) {
 	if !isStructOrStructPointer(t) {
-		return nil, xerrors.Errorf("typeToTableHeaders called with a non-struct or a non-pointer-to-a-struct type")
+		return nil, "", xerrors.Errorf("typeToTableHeaders called with a non-struct or a non-pointer-to-a-struct type")
 	}
 	if t.Kind() == reflect.Pointer {
 		t = t.Elem()
 	}
 
 	headers := []string{}
+	defaultSortName := ""
 	for i := 0; i < t.NumField(); i++ {
 		field := t.Field(i)
-		name, recursive, err := parseTableStructTag(field)
+		name, defaultSort, recursive, err := parseTableStructTag(field)
 		if err != nil {
-			return nil, xerrors.Errorf("parse struct tags for field %q in type %q: %w", field.Name, t.String(), err)
+			return nil, "", xerrors.Errorf("parse struct tags for field %q in type %q: %w", field.Name, t.String(), err)
 		}
 		if name == "" {
 			continue
 		}
+		if defaultSort {
+			if defaultSortName != "" {
+				return nil, "", xerrors.Errorf("multiple fields marked as default sort in type %q", t.String())
+			}
+			defaultSortName = name
+		}
 
 		fieldType := field.Type
 		if recursive {
 			if !isStructOrStructPointer(fieldType) {
-				return nil, xerrors.Errorf("field %q in type %q is marked as recursive but does not contain a struct or a pointer to a struct", field.Name, t.String())
+				return nil, "", xerrors.Errorf("field %q in type %q is marked as recursive but does not contain a struct or a pointer to a struct", field.Name, t.String())
 			}
 
-			childNames, err := typeToTableHeaders(fieldType)
+			childNames, _, err := typeToTableHeaders(fieldType)
 			if err != nil {
-				return nil, xerrors.Errorf("get child field header names for field %q in type %q: %w", field.Name, fieldType.String(), err)
+				return nil, "", xerrors.Errorf("get child field header names for field %q in type %q: %w", field.Name, fieldType.String(), err)
 			}
 			for _, childName := range childNames {
 				headers = append(headers, fmt.Sprintf("%s %s", name, childName))
@@ -252,7 +268,11 @@ func typeToTableHeaders(t reflect.Type) ([]string, error) {
 		headers = append(headers, name)
 	}
 
-	return headers, nil
+	if defaultSortName == "" {
+		return nil, "", xerrors.Errorf("no field marked as default_sort in type %q", t.String())
+	}
+
+	return headers, defaultSortName, nil
 }
 
 // valueToTableMap converts a struct to a map of column name to value. If the
@@ -276,7 +296,7 @@ func valueToTableMap(val reflect.Value) (map[string]any, error) {
 	for i := 0; i < val.NumField(); i++ {
 		field := val.Type().Field(i)
 		fieldVal := val.Field(i)
-		name, recursive, err := parseTableStructTag(field)
+		name, _, recursive, err := parseTableStructTag(field)
 		if err != nil {
 			return nil, xerrors.Errorf("parse struct tags for field %q in type %T: %w", field.Name, val, err)
 		}
@@ -309,18 +329,3 @@ func valueToTableMap(val reflect.Value) (map[string]any, error) {
 
 	return row, nil
 }
-
-// TableHeaders returns the table header names of all
-// fields in tSlice. tSlice must be a slice of some type.
-func TableHeaders(tSlice any) ([]string, error) {
-	v := reflect.Indirect(reflect.ValueOf(tSlice))
-	rawHeaders, err := typeToTableHeaders(v.Type().Elem())
-	if err != nil {
-		return nil, xerrors.Errorf("type to table headers: %w", err)
-	}
-	out := make([]string, 0, len(rawHeaders))
-	for _, hdr := range rawHeaders {
-		out = append(out, strings.Replace(hdr, " ", "_", -1))
-	}
-	return out, nil
-}
diff --git a/cli/list.go b/cli/list.go
@@ -131,10 +131,8 @@ func list() *cobra.Command {
 		},
 	}
 
-	availColumns, err := cliui.TableHeaders(displayWorkspaces)
-	if err != nil {
-		panic(err)
-	}
+	// TODO: fix this
+	availColumns := []string{"name"}
 	columnString := strings.Join(availColumns[:], ", ")
 
 	cmd.Flags().BoolVarP(&all, "all", "a", false,