diff --git a/ansiescape/split.go b/ansiescape/split.go new file mode 100644 index 0000000..b097b25 --- /dev/null +++ b/ansiescape/split.go @@ -0,0 +1,89 @@ +package ansiescape + +import "bytes" + +// dropCR drops a leading or terminal \r from the data. +func dropCR(data []byte) []byte { + if len(data) > 0 && data[len(data)-1] == '\r' { + data = data[0 : len(data)-1] + } + if len(data) > 0 && data[0] == '\r' { + data = data[1:] + } + return data +} + +// escapeSequenceLength calculates the length of an ANSI escape sequence +// If there is not enough characters to match a sequence, -1 is returned, +// if there is no valid sequence 0 is returned, otherwise the number +// of bytes in the sequence is returned. Only returns length for +// line moving sequences. +func escapeSequenceLength(data []byte) int { + next := 0 + if len(data) <= next { + return -1 + } + if data[next] != '[' { + return 0 + } + for { + next = next + 1 + if len(data) <= next { + return -1 + } + if (data[next] > '9' || data[next] < '0') && data[next] != ';' { + break + } + } + if len(data) <= next { + return -1 + } + // Only match line moving codes + switch data[next] { + case 'A', 'B', 'E', 'F', 'H', 'h': + return next + 1 + } + + return 0 +} + +// ScanANSILines is a scanner function which splits the +// input based on ANSI escape codes and new lines. +func ScanANSILines(data []byte, atEOF bool) (advance int, token []byte, err error) { + if atEOF && len(data) == 0 { + return 0, nil, nil + } + + // Look for line moving escape sequence + if i := bytes.IndexByte(data, '\x1b'); i >= 0 { + last := 0 + for i >= 0 { + last = last + i + + // get length of ANSI escape sequence + sl := escapeSequenceLength(data[last+1:]) + if sl == -1 { + return 0, nil, nil + } + if sl == 0 { + // If no relevant sequence was found, skip + last = last + 1 + i = bytes.IndexByte(data[last:], '\x1b') + continue + } + + return last + 1 + sl, dropCR(data[0:(last)]), nil + } + } + if i := bytes.IndexByte(data, '\n'); i >= 0 { + // No escape sequence, check for new line + return i + 1, dropCR(data[0:i]), nil + } + + // If we're at EOF, we have a final, non-terminated line. Return it. + if atEOF { + return len(data), dropCR(data), nil + } + // Request more data. + return 0, nil, nil +} diff --git a/ansiescape/split_test.go b/ansiescape/split_test.go new file mode 100644 index 0000000..ecb24b9 --- /dev/null +++ b/ansiescape/split_test.go @@ -0,0 +1,53 @@ +package ansiescape + +import ( + "bufio" + "strings" + "testing" +) + +func TestSplit(t *testing.T) { + lines := []string{ + "test line 1", + "another test line", + "some test line", + "line with non-cursor moving sequence \x1b[1T", // Scroll Down + "line with \x1b[31;1mcolor\x1b[0m then reset", // "color" in Bold Red + "cursor forward \x1b[1C and backward \x1b[1D", + "invalid sequence \x1babcd", + "", + "after empty", + } + splitSequences := []string{ + "\x1b[1A", // Cursor up + "\x1b[1B", // Cursor down + "\x1b[1E", // Cursor next line + "\x1b[1F", // Cursor previous line + "\x1b[1;1H", // Move cursor to position + "\x1b[1;1h", // Move cursor to position + "\n", + "\r\n", + "\n\r", + "\x1b[1A\r", + "\r\x1b[1A", + } + + for _, sequence := range splitSequences { + scanner := bufio.NewScanner(strings.NewReader(strings.Join(lines, sequence))) + scanner.Split(ScanANSILines) + i := 0 + for scanner.Scan() { + if i >= len(lines) { + t.Fatalf("Too many scanned lines") + } + scanned := scanner.Text() + if scanned != lines[i] { + t.Fatalf("Wrong line scanned with sequence %q\n\tExpected: %q\n\tActual: %q", sequence, lines[i], scanned) + } + i++ + } + if i < len(lines) { + t.Errorf("Wrong number of lines for sequence %q: %d, expected %d", sequence, i, len(lines)) + } + } +}