Move ansi stripper to util package

2 files changed, 88 insertions, 0 deletions

diff --git a/util/util.go b/util/util.go
new file mode 100644
index 0000000..deb528b
--- /dev/null
+++ b/util/util.go
@@ -0,0 +1,38 @@
+package util
+
+import (
+	"io"
+	"regexp"
+)
+
+// Strip ansi control characters from source
+//
+// The source reader is automatically closed when stripper reader is
+type AnsiStripper struct {
+	regex  regexp.Regexp
+	source io.ReadCloser
+}
+
+func NewAnsiStripper(source io.ReadCloser) *AnsiStripper {
+	striper := new(AnsiStripper)
+	striper.regex = *regexp.MustCompile(`\x1b[[0-9;]*m`)
+	striper.source = source
+
+	return striper
+}
+
+func (rc AnsiStripper) Read(p []byte) (int, error) {
+	n, err := rc.source.Read(p)
+	if err != io.EOF && err != nil {
+		return n, err
+	}
+
+	stripped := rc.regex.ReplaceAll(p[:n], []byte(""))
+	n = copy(p, stripped)
+
+	return n, err
+}
+
+func (rc AnsiStripper) Close() error {
+	return rc.source.Close()
+}
diff --git a/util/util_test.go b/util/util_test.go
new file mode 100644
index 0000000..1c4ae39
--- /dev/null
+++ b/util/util_test.go
@@ -0,0 +1,50 @@
+package util_test;
+
+import (
+	"io"
+	"nonsense-time/util"
+	"strings"
+	"testing"
+)
+
+type stringReadCloser struct {
+	*strings.Reader
+}
+
+func newStringReadCloser(s string) stringReadCloser {
+	rc := stringReadCloser{}
+	rc.Reader = strings.NewReader(s)
+
+	return rc
+}
+
+func (stringReadCloser) Close() error {
+	return nil
+}
+
+func testAnsiStripper(t *testing.T, input string, expected string) {
+	reader := newStringReadCloser(input)
+	cleanReader := util.NewAnsiStripper(reader)
+	defer cleanReader.Close()
+
+	buf := new(strings.Builder)
+	n, err := io.Copy(buf, cleanReader)
+	if err != nil {
+		t.Fatal("Error while copying cleaned text to output", err)
+	}
+
+	result := buf.String()
+
+	if n != int64(len(expected)) {
+		t.Errorf("Expected to write %d characters but wrote %d\n", 3, n)
+	}
+
+	if result != expected {
+		t.Errorf("Expected string `%s` but wrote `%s`", "abc", result)
+	}
+}
+
+func TestStripAnsiColors(t *testing.T) {
+	testAnsiStripper(t, "a\x1b[31mbc", "abc")
+	testAnsiStripper(t, "[\x1b[32minfo\x1b[39m]", "[info]")
+}