2020-09-20 00:00:50 +08:00
|
|
|
package ansi
|
|
|
|
|
|
|
|
import (
|
2021-08-30 23:18:50 +08:00
|
|
|
"bytes"
|
2020-09-20 00:00:50 +08:00
|
|
|
"io"
|
2021-08-30 23:18:50 +08:00
|
|
|
"unicode/utf8"
|
2020-09-20 00:00:50 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
type Writer struct {
|
|
|
|
Forward io.Writer
|
|
|
|
|
|
|
|
ansi bool
|
2021-08-30 23:18:50 +08:00
|
|
|
ansiseq bytes.Buffer
|
|
|
|
lastseq bytes.Buffer
|
2020-09-20 00:00:50 +08:00
|
|
|
seqchanged bool
|
2021-08-30 23:18:50 +08:00
|
|
|
runeBuf []byte
|
2020-09-20 00:00:50 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// Write is used to write content to the ANSI buffer.
|
|
|
|
func (w *Writer) Write(b []byte) (int, error) {
|
|
|
|
for _, c := range string(b) {
|
2021-08-30 23:18:50 +08:00
|
|
|
if c == Marker {
|
2020-09-20 00:00:50 +08:00
|
|
|
// ANSI escape sequence
|
|
|
|
w.ansi = true
|
|
|
|
w.seqchanged = true
|
2021-08-30 23:18:50 +08:00
|
|
|
_, _ = w.ansiseq.WriteRune(c)
|
2020-09-20 00:00:50 +08:00
|
|
|
} else if w.ansi {
|
2021-08-30 23:18:50 +08:00
|
|
|
_, _ = w.ansiseq.WriteRune(c)
|
|
|
|
if IsTerminator(c) {
|
2020-09-20 00:00:50 +08:00
|
|
|
// ANSI sequence terminated
|
|
|
|
w.ansi = false
|
|
|
|
|
2021-08-30 23:18:50 +08:00
|
|
|
if bytes.HasSuffix(w.ansiseq.Bytes(), []byte("[0m")) {
|
2020-09-20 00:00:50 +08:00
|
|
|
// reset sequence
|
2021-08-30 23:18:50 +08:00
|
|
|
w.lastseq.Reset()
|
|
|
|
w.seqchanged = false
|
|
|
|
} else if c == 'm' {
|
2020-09-20 00:00:50 +08:00
|
|
|
// color code
|
2021-08-30 23:18:50 +08:00
|
|
|
_, _ = w.lastseq.Write(w.ansiseq.Bytes())
|
2020-09-20 00:00:50 +08:00
|
|
|
}
|
2021-08-30 23:18:50 +08:00
|
|
|
|
|
|
|
_, _ = w.ansiseq.WriteTo(w.Forward)
|
2020-09-20 00:00:50 +08:00
|
|
|
}
|
|
|
|
} else {
|
2021-08-30 23:18:50 +08:00
|
|
|
_, err := w.writeRune(c)
|
2020-09-20 00:00:50 +08:00
|
|
|
if err != nil {
|
|
|
|
return 0, err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
return len(b), nil
|
|
|
|
}
|
|
|
|
|
2021-08-30 23:18:50 +08:00
|
|
|
func (w *Writer) writeRune(r rune) (int, error) {
|
|
|
|
if w.runeBuf == nil {
|
|
|
|
w.runeBuf = make([]byte, utf8.UTFMax)
|
|
|
|
}
|
|
|
|
n := utf8.EncodeRune(w.runeBuf, r)
|
|
|
|
return w.Forward.Write(w.runeBuf[:n])
|
|
|
|
}
|
|
|
|
|
2020-09-20 00:00:50 +08:00
|
|
|
func (w *Writer) LastSequence() string {
|
2021-08-30 23:18:50 +08:00
|
|
|
return w.lastseq.String()
|
2020-09-20 00:00:50 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) ResetAnsi() {
|
|
|
|
if !w.seqchanged {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
_, _ = w.Forward.Write([]byte("\x1b[0m"))
|
|
|
|
}
|
|
|
|
|
|
|
|
func (w *Writer) RestoreAnsi() {
|
2021-08-30 23:18:50 +08:00
|
|
|
_, _ = w.Forward.Write(w.lastseq.Bytes())
|
2020-09-20 00:00:50 +08:00
|
|
|
}
|