status.go

package termstatus

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"strings"

	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/text/width"
)

// Terminal is used to write messages and display status lines which can be
// updated. When the output is redirected to a file, the status lines are not
// printed.
type Terminal struct {
	wr              *bufio.Writer
	fd              uintptr
	errWriter       io.Writer
	buf             *bytes.Buffer
	msg             chan message
	status          chan status
	canUpdateStatus bool

	// will be closed when the goroutine which runs Run() terminates, so it'll
	// yield a default value immediately
	closed chan struct{}

	clearCurrentLine func(io.Writer, uintptr)
	moveCursorUp     func(io.Writer, uintptr, int)
}

type message struct {
	line string
	err  bool
}

type status struct {
	lines []string
}

type fder interface {
	Fd() uintptr
}

// New returns a new Terminal for wr. A goroutine is started to update the
// terminal. It is terminated when ctx is cancelled. When wr is redirected to
// a file (e.g. via shell output redirection) or is just an io.Writer (not the
// open *os.File for stdout), no status lines are printed. The status lines and
// normal output (via Print/Printf) are written to wr, error messages are
// written to errWriter. If disableStatus is set to true, no status messages
// are printed even if the terminal supports it.
func New(wr io.Writer, errWriter io.Writer, disableStatus bool) *Terminal {
	t := &Terminal{
		wr:        bufio.NewWriter(wr),
		errWriter: errWriter,
		buf:       bytes.NewBuffer(nil),
		msg:       make(chan message),
		status:    make(chan status),
		closed:    make(chan struct{}),
	}

	if disableStatus {
		return t
	}

	if d, ok := wr.(fder); ok && canUpdateStatus(d.Fd()) {
		// only use the fancy status code when we're running on a real terminal.
		t.canUpdateStatus = true
		t.fd = d.Fd()
		t.clearCurrentLine = clearCurrentLine(wr, t.fd)
		t.moveCursorUp = moveCursorUp(wr, t.fd)
	}

	return t
}

// CanUpdateStatus return whether the status output is updated in place.
func (t *Terminal) CanUpdateStatus() bool {
	return t.canUpdateStatus
}

// Run updates the screen. It should be run in a separate goroutine. When
// ctx is cancelled, the status lines are cleanly removed.
func (t *Terminal) Run(ctx context.Context) {
	defer close(t.closed)
	if t.canUpdateStatus {
		t.run(ctx)
		return
	}

	t.runWithoutStatus(ctx)
}

// run listens on the channels and updates the terminal screen.
func (t *Terminal) run(ctx context.Context) {
	var status []string
	for {
		select {
		case <-ctx.Done():
			if !IsProcessBackground(t.fd) {
				t.undoStatus(len(status))
			}

			return

		case msg := <-t.msg:
			if IsProcessBackground(t.fd) {
				// ignore all messages, do nothing, we are in the background process group
				continue
			}
			t.clearCurrentLine(t.wr, t.fd)

			var dst io.Writer
			if msg.err {
				dst = t.errWriter

				// assume t.wr and t.errWriter are different, so we need to
				// flush clearing the current line
				err := t.wr.Flush()
				if err != nil {
					fmt.Fprintf(os.Stderr, "flush failed: %v\n", err)
				}
			} else {
				dst = t.wr
			}

			if _, err := io.WriteString(dst, msg.line); err != nil {
				fmt.Fprintf(os.Stderr, "write failed: %v\n", err)
				continue
			}

			t.writeStatus(status)

			if err := t.wr.Flush(); err != nil {
				fmt.Fprintf(os.Stderr, "flush failed: %v\n", err)
			}

		case stat := <-t.status:
			if IsProcessBackground(t.fd) {
				// ignore all messages, do nothing, we are in the background process group
				continue
			}

			status = status[:0]
			status = append(status, stat.lines...)
			t.writeStatus(status)
		}
	}
}

func (t *Terminal) writeStatus(status []string) {
	for _, line := range status {
		t.clearCurrentLine(t.wr, t.fd)

		_, err := t.wr.WriteString(line)
		if err != nil {
			fmt.Fprintf(os.Stderr, "write failed: %v\n", err)
		}

		// flush is needed so that the current line is updated
		err = t.wr.Flush()
		if err != nil {
			fmt.Fprintf(os.Stderr, "flush failed: %v\n", err)
		}
	}

	if len(status) > 0 {
		t.moveCursorUp(t.wr, t.fd, len(status)-1)
	}

	err := t.wr.Flush()
	if err != nil {
		fmt.Fprintf(os.Stderr, "flush failed: %v\n", err)
	}
}

// runWithoutStatus listens on the channels and just prints out the messages,
// without status lines.
func (t *Terminal) runWithoutStatus(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case msg := <-t.msg:
			var flush func() error

			var dst io.Writer
			if msg.err {
				dst = t.errWriter
			} else {
				dst = t.wr
				flush = t.wr.Flush
			}

			if _, err := io.WriteString(dst, msg.line); err != nil {
				fmt.Fprintf(os.Stderr, "write failed: %v\n", err)
			}

			if flush == nil {
				continue
			}

			if err := flush(); err != nil {
				fmt.Fprintf(os.Stderr, "flush failed: %v\n", err)
			}

		case stat := <-t.status:
			for _, line := range stat.lines {
				// Ensure that each message ends with exactly one newline.
				fmt.Fprintln(t.wr, strings.TrimRight(line, "\n"))
			}
			if err := t.wr.Flush(); err != nil {
				fmt.Fprintf(os.Stderr, "flush failed: %v\n", err)
			}
		}
	}
}

func (t *Terminal) undoStatus(lines int) {
	for i := 0; i < lines; i++ {
		t.clearCurrentLine(t.wr, t.fd)

		_, err := t.wr.WriteRune('\n')
		if err != nil {
			fmt.Fprintf(os.Stderr, "write failed: %v\n", err)
		}

		// flush is needed so that the current line is updated
		err = t.wr.Flush()
		if err != nil {
			fmt.Fprintf(os.Stderr, "flush failed: %v\n", err)
		}
	}

	t.moveCursorUp(t.wr, t.fd, lines)

	err := t.wr.Flush()
	if err != nil {
		fmt.Fprintf(os.Stderr, "flush failed: %v\n", err)
	}
}

func (t *Terminal) print(line string, isErr bool) {
	// make sure the line ends with a line break
	if line[len(line)-1] != '\n' {
		line += "\n"
	}

	select {
	case t.msg <- message{line: line, err: isErr}:
	case <-t.closed:
	}
}

// Print writes a line to the terminal.
func (t *Terminal) Print(line string) {
	t.print(line, false)
}

// Printf uses fmt.Sprintf to write a line to the terminal.
func (t *Terminal) Printf(msg string, args ...interface{}) {
	s := fmt.Sprintf(msg, args...)
	t.Print(s)
}

// Error writes an error to the terminal.
func (t *Terminal) Error(line string) {
	t.print(line, true)
}

// Errorf uses fmt.Sprintf to write an error line to the terminal.
func (t *Terminal) Errorf(msg string, args ...interface{}) {
	s := fmt.Sprintf(msg, args...)
	t.Error(s)
}

// Truncate s to fit in width (number of terminal cells) w.
// If w is negative, returns the empty string.
func truncate(s string, w int) string {
	if len(s) < w {
		// Since the display width of a character is at most 2
		// and all of ASCII (single byte per rune) has width 1,
		// no character takes more bytes to encode than its width.
		return s
	}

	for i, r := range s {
		// Determine width of the rune. This cannot be determined without
		// knowing the terminal font, so let's just be careful and treat
		// all ambigous characters as full-width, i.e., two cells.
		wr := 2
		switch width.LookupRune(r).Kind() {
		case width.Neutral, width.EastAsianNarrow:
			wr = 1
		}

		w -= wr
		if w < 0 {
			return s[:i]
		}
	}

	return s
}

// SetStatus updates the status lines.
func (t *Terminal) SetStatus(lines []string) {
	if len(lines) == 0 {
		return
	}

	// only truncate interactive status output
	var width int
	if t.canUpdateStatus {
		var err error
		width, _, err = terminal.GetSize(int(t.fd))
		if err != nil || width <= 0 {
			// use 80 columns by default
			width = 80
		}
	}

	// make sure that all lines have a line break and are not too long
	for i, line := range lines {
		line = strings.TrimRight(line, "\n")
		if width > 0 {
			line = truncate(line, width-2)
		}
		lines[i] = line + "\n"
	}

	// make sure the last line does not have a line break
	last := len(lines) - 1
	lines[last] = strings.TrimRight(lines[last], "\n")

	select {
	case t.status <- status{lines: lines}:
	case <-t.closed:
	}
}