table.go

/*
Copyright 2017 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package asciitable implements a simple ASCII table formatter for printing
// tabular values into a text terminal.
package asciitable

import (
	"bytes"
	"fmt"
	"strings"
	"text/tabwriter"
)

// column represents a column in the table. Contains the maximum width of the
// column as well as the title.
type column struct {
	width int
	title string
}

// Table holds tabular values in a rows and columns format.
type Table struct {
	columns []column
	rows    [][]string
}

// MakeTable creates a new instance of the table with given column names.
func MakeTable(headers []string) Table {
	t := MakeHeadlessTable(len(headers))
	for i := range t.columns {
		t.columns[i].title = headers[i]
		t.columns[i].width = len(headers[i])
	}
	return t
}

// MakeTable creates a new instance of the table without any column names.
// The number of columns is required.
func MakeHeadlessTable(columnCount int) Table {
	return Table{
		columns: make([]column, columnCount),
		rows:    make([][]string, 0),
	}
}

// AddRow adds a row of cells to the table.
func (t *Table) AddRow(row []string) {
	limit := min(len(row), len(t.columns))
	for i := 0; i < limit; i++ {
		cellWidth := len(row[i])
		t.columns[i].width = max(cellWidth, t.columns[i].width)
	}
	t.rows = append(t.rows, row[:limit])
}

// AsBuffer returns a *bytes.Buffer with the printed output of the table.
func (t *Table) AsBuffer() *bytes.Buffer {
	var buffer bytes.Buffer

	writer := tabwriter.NewWriter(&buffer, 5, 0, 1, ' ', 0)
	template := strings.Repeat("%v\t", len(t.columns))

	// Header and separator.
	if !t.IsHeadless() {
		var colh []interface{}
		var cols []interface{}

		for _, col := range t.columns {
			colh = append(colh, col.title)
			cols = append(cols, strings.Repeat("-", col.width))
		}
		fmt.Fprintf(writer, template+"\n", colh...)
		fmt.Fprintf(writer, template+"\n", cols...)
	}

	// Body.
	for _, row := range t.rows {
		var rowi []interface{}
		for _, cell := range row {
			rowi = append(rowi, cell)
		}
		fmt.Fprintf(writer, template+"\n", rowi...)
	}

	writer.Flush()
	return &buffer
}

// IsHeadless returns true if none of the table title cells contains any text.
func (t *Table) IsHeadless() bool {
	total := 0
	for i := range t.columns {
		total += len(t.columns[i].title)
	}
	return total == 0
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}