121 lines
2.6 KiB
Go
121 lines
2.6 KiB
Go
package truncate
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
|
|
"github.com/mattn/go-runewidth"
|
|
|
|
"github.com/muesli/reflow/ansi"
|
|
)
|
|
|
|
type Writer struct {
|
|
width uint
|
|
tail string
|
|
|
|
ansiWriter *ansi.Writer
|
|
buf bytes.Buffer
|
|
ansi bool
|
|
}
|
|
|
|
func NewWriter(width uint, tail string) *Writer {
|
|
w := &Writer{
|
|
width: width,
|
|
tail: tail,
|
|
}
|
|
w.ansiWriter = &ansi.Writer{
|
|
Forward: &w.buf,
|
|
}
|
|
return w
|
|
}
|
|
|
|
func NewWriterPipe(forward io.Writer, width uint, tail string) *Writer {
|
|
return &Writer{
|
|
width: width,
|
|
tail: tail,
|
|
ansiWriter: &ansi.Writer{
|
|
Forward: forward,
|
|
},
|
|
}
|
|
}
|
|
|
|
// Bytes is shorthand for declaring a new default truncate-writer instance,
|
|
// used to immediately truncate a byte slice.
|
|
func Bytes(b []byte, width uint) []byte {
|
|
return BytesWithTail(b, width, []byte(""))
|
|
}
|
|
|
|
// Bytes is shorthand for declaring a new default truncate-writer instance,
|
|
// used to immediately truncate a byte slice. A tail is then added to the
|
|
// end of the byte slice.
|
|
func BytesWithTail(b []byte, width uint, tail []byte) []byte {
|
|
f := NewWriter(width, string(tail))
|
|
_, _ = f.Write(b)
|
|
|
|
return f.Bytes()
|
|
}
|
|
|
|
// String is shorthand for declaring a new default truncate-writer instance,
|
|
// used to immediately truncate a string.
|
|
func String(s string, width uint) string {
|
|
return StringWithTail(s, width, "")
|
|
}
|
|
|
|
// StringWithTail is shorthand for declaring a new default truncate-writer instance,
|
|
// used to immediately truncate a string. A tail is then added to the end of the
|
|
// string.
|
|
func StringWithTail(s string, width uint, tail string) string {
|
|
return string(BytesWithTail([]byte(s), width, []byte(tail)))
|
|
}
|
|
|
|
// Write truncates content at the given printable cell width, leaving any
|
|
// ansi sequences intact.
|
|
func (w *Writer) Write(b []byte) (int, error) {
|
|
tw := ansi.PrintableRuneWidth(w.tail)
|
|
if w.width < uint(tw) {
|
|
return w.buf.WriteString(w.tail)
|
|
}
|
|
|
|
w.width -= uint(tw)
|
|
var curWidth uint
|
|
|
|
for _, c := range string(b) {
|
|
if c == ansi.Marker {
|
|
// ANSI escape sequence
|
|
w.ansi = true
|
|
} else if w.ansi {
|
|
if ansi.IsTerminator(c) {
|
|
// ANSI sequence terminated
|
|
w.ansi = false
|
|
}
|
|
} else {
|
|
curWidth += uint(runewidth.RuneWidth(c))
|
|
}
|
|
|
|
if curWidth > w.width {
|
|
n, err := w.buf.WriteString(w.tail)
|
|
if w.ansiWriter.LastSequence() != "" {
|
|
w.ansiWriter.ResetAnsi()
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
_, err := w.ansiWriter.Write([]byte(string(c)))
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
|
|
return len(b), nil
|
|
}
|
|
|
|
// Bytes returns the truncated result as a byte slice.
|
|
func (w *Writer) Bytes() []byte {
|
|
return w.buf.Bytes()
|
|
}
|
|
|
|
// String returns the truncated result as a string.
|
|
func (w *Writer) String() string {
|
|
return w.buf.String()
|
|
}
|