terminal_writer.go

package term

import (
	"io"
	"os"

	"github.com/docker/docker/pkg/term"
	wordwrap "github.com/mitchellh/go-wordwrap"
	kterm "k8s.io/kubernetes/pkg/kubectl/util/term"
)

type wordWrapWriter struct {
	limit  uint
	writer io.Writer
}

// NewResponsiveWriter creates a Writer that detects the column width of the
// terminal we are in, and adjusts every line width to fit and use recommended
// terminal sizes for better readability. Does proper word wrapping automatically.
//    if terminal width >= 120 columns		use 120 columns
//    if terminal width >= 100 columns		use 100 columns
//    if terminal width >=  80 columns		use  80 columns
// In case we're not in a terminal or if it's smaller than 80 columns width,
// doesn't do any wrapping.
func NewResponsiveWriter(w io.Writer) io.Writer {
	file, ok := w.(*os.File)
	if !ok {
		return w
	}
	fd := file.Fd()
	if !term.IsTerminal(fd) {
		return w
	}

	terminalSize := kterm.GetSize(fd)
	if terminalSize == nil {
		return w
	}

	var limit uint
	switch {
	case terminalSize.Width >= 120:
		limit = 120
	case terminalSize.Width >= 100:
		limit = 100
	case terminalSize.Width >= 80:
		limit = 80
	}

	return NewWordWrapWriter(w, limit)
}

// NewWordWrapWriter is a Writer that supports a limit of characters on every line
// and does auto word wrapping that respects that limit.
func NewWordWrapWriter(w io.Writer, limit uint) io.Writer {
	return &wordWrapWriter{
		limit:  limit,
		writer: w,
	}
}

func (w wordWrapWriter) Write(p []byte) (nn int, err error) {
	if w.limit == 0 {
		return w.writer.Write(p)
	}
	original := string(p)
	wrapped := wordwrap.WrapString(original, w.limit)
	l, e := w.writer.Write([]byte(wrapped))

	// This wrapped word writer could modify input in some cases such that
	// the length of output is not equal to the length of input and this
	// can cause ErrShortWrite error. So to avoid the error, make sure the
	// returned length of output bytes is equal to the length of input.
	// This returns length of original in non-erroneous cases, and in erroneous
	// cases min(len(original), l).
	if e == nil || l > len(original) {
		l = len(original)
	}

	return l, e
}

// NewPunchCardWriter is a NewWordWrapWriter that limits the line width to 80 columns.
func NewPunchCardWriter(w io.Writer) io.Writer {
	return NewWordWrapWriter(w, 80)
}

type maxWidthWriter struct {
	maxWidth     uint
	currentWidth uint
	written      uint
	writer       io.Writer
}

// NewMaxWidthWriter is a Writer that supports a limit of characters on every
// line, but doesn't do any word wrapping automatically.
func NewMaxWidthWriter(w io.Writer, maxWidth uint) io.Writer {
	return &maxWidthWriter{
		maxWidth: maxWidth,
		writer:   w,
	}
}

func (m maxWidthWriter) Write(p []byte) (nn int, err error) {
	for _, b := range p {
		if m.currentWidth == m.maxWidth {
			m.writer.Write([]byte{'\n'})
			m.currentWidth = 0
		}
		if b == '\n' {
			m.currentWidth = 0
		}
		_, err := m.writer.Write([]byte{b})
		if err != nil {
			return int(m.written), err
		}
		m.written++
		m.currentWidth++
	}
	return len(p), nil
}