padding.go

/// A modified version of padding from reflow, with support of iTerm2 proprietary image sequences.
/// The code is modified so it will not insert \x1b]0m in an image sequence.

package padding

import (
	"bytes"
	"io"
	"strings"

	"github.com/mattn/go-runewidth"
	"github.com/muesli/reflow/ansi"
)

type PaddingFunc func(w io.Writer)

type Writer struct {
	Padding uint
	PadFunc PaddingFunc

	ansiWriter         *ansi.Writer
	buf                bytes.Buffer
	cache              bytes.Buffer
	lineLen            int
	ansi               bool
	iterm              bool
	ansiWaitSecondChar bool
}

func NewWriter(width uint, paddingFunc PaddingFunc) *Writer {
	w := &Writer{
		Padding: width,
		PadFunc: paddingFunc,
	}
	w.ansiWriter = &ansi.Writer{
		Forward: &w.buf,
	}
	return w
}

func NewWriterPipe(forward io.Writer, width uint, paddingFunc PaddingFunc) *Writer {
	return &Writer{
		Padding: width,
		PadFunc: paddingFunc,
		ansiWriter: &ansi.Writer{
			Forward: forward,
		},
	}
}

// Bytes is shorthand for declaring a new default padding-writer instance,
// used to immediately pad a byte slice.
func Bytes(b []byte, width uint) []byte {
	f := NewWriter(width, nil)
	_, _ = f.Write(b)
	_ = f.Flush()

	return f.Bytes()
}

// String is shorthand for declaring a new default padding-writer instance,
// used to immediately pad a string.
func String(s string, width uint) string {
	return string(Bytes([]byte(s), width))
}

// Write is used to write content to the padding buffer.
func (w *Writer) Write(b []byte) (int, error) {
	for _, c := range string(b) {
		if c == '\x1B' {
			// ANSI escape sequence
			w.ansi = true
			w.ansiWaitSecondChar = true
		} else if w.ansi {
			if w.ansiWaitSecondChar {
				w.iterm = c == ']'
				w.ansiWaitSecondChar = false
			}

			if !w.iterm {
				if (c >= 0x40 && c <= 0x5a) || (c >= 0x61 && c <= 0x7a) {
					// ANSI sequence terminated
					w.ansi = false
					w.ansiWaitSecondChar = false
					w.iterm = false
				}
			} else {
				if c == '\a' {
					// ANSI sequence terminated
					w.ansi = false
					w.ansiWaitSecondChar = false
					w.iterm = false
				}
			}
		} else {
			w.lineLen += runewidth.StringWidth(string(c))

			if c == '\n' {
				// end of current line
				err := w.pad()
				if err != nil {
					return 0, err
				}
				w.ansiWriter.ResetAnsi()
				w.lineLen = 0
			}
		}

		_, err := w.ansiWriter.Write([]byte(string(c)))
		if err != nil {
			return 0, err
		}
	}

	return len(b), nil
}

func (w *Writer) pad() error {
	if w.Padding > 0 && uint(w.lineLen) < w.Padding {
		if w.PadFunc != nil {
			for i := 0; i < int(w.Padding)-w.lineLen; i++ {
				w.PadFunc(w.ansiWriter)
			}
		} else {
			_, err := w.ansiWriter.Write([]byte(strings.Repeat(" ", int(w.Padding)-w.lineLen)))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Close will finish the padding operation.
func (w *Writer) Close() (err error) {
	return w.Flush()
}

// Bytes returns the padded result as a byte slice.
func (w *Writer) Bytes() []byte {
	return w.cache.Bytes()
}

// String returns the padded result as a string.
func (w *Writer) String() string {
	return w.cache.String()
}

// Flush will finish the padding operation. Always call it before trying to
// retrieve the final result.
func (w *Writer) Flush() (err error) {
	if w.lineLen != 0 {
		if err = w.pad(); err != nil {
			return
		}
	}

	w.cache.Reset()
	_, err = w.buf.WriteTo(&w.cache)
	w.lineLen = 0
	w.ansi = false

	return
}