Add ANSI compatible io.Writer

ansi/writer.go

@@ -0,0 +1,65 @@
+package ansi
+
+import (
+	"io"
+	"strings"
+)
+
+type Writer struct {
+	Forward io.Writer
+
+	ansi       bool
+	ansiseq    string
+	lastseq    string
+	seqchanged bool
+}
+
+// Write is used to write content to the ANSI buffer.
+func (w *Writer) Write(b []byte) (int, error) {
+	for _, c := range string(b) {
+		if c == '\x1B' {
+			// ANSI escape sequence
+			w.ansi = true
+			w.seqchanged = true
+			w.ansiseq += string(c)
+		} else if w.ansi {
+			w.ansiseq += string(c)
+			if (c >= 0x41 && c <= 0x5a) || (c >= 0x61 && c <= 0x7a) {
+				// ANSI sequence terminated
+				w.ansi = false
+
+				_, _ = w.Forward.Write([]byte(w.ansiseq))
+				if strings.HasSuffix(w.ansiseq, "[0m") {
+					// reset sequence
+					w.lastseq = ""
+				} else if strings.HasSuffix(w.ansiseq, "m") {
+					// color code
+					w.lastseq = w.ansiseq
+				}
+				w.ansiseq = ""
+			}
+		} else {
+			_, err := w.Forward.Write([]byte(string(c)))
+			if err != nil {
+				return 0, err
+			}
+		}
+	}
+
+	return len(b), nil
+}
+
+func (w *Writer) LastSequence() string {
+	return w.lastseq
+}
+
+func (w *Writer) ResetAnsi() {
+	if !w.seqchanged {
+		return
+	}
+	_, _ = w.Forward.Write([]byte("\x1b[0m"))
+}
+
+func (w *Writer) RestoreAnsi() {
+	_, _ = w.Forward.Write([]byte(w.lastseq))
+}