writer.go 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576
  1. package ansi
  2. import (
  3. "bytes"
  4. "io"
  5. "unicode/utf8"
  6. )
  7. type Writer struct {
  8. Forward io.Writer
  9. ansi bool
  10. ansiseq bytes.Buffer
  11. lastseq bytes.Buffer
  12. seqchanged bool
  13. runeBuf []byte
  14. }
  15. // Write is used to write content to the ANSI buffer.
  16. func (w *Writer) Write(b []byte) (int, error) {
  17. for _, c := range string(b) {
  18. if c == Marker {
  19. // ANSI escape sequence
  20. w.ansi = true
  21. w.seqchanged = true
  22. _, _ = w.ansiseq.WriteRune(c)
  23. } else if w.ansi {
  24. _, _ = w.ansiseq.WriteRune(c)
  25. if IsTerminator(c) {
  26. // ANSI sequence terminated
  27. w.ansi = false
  28. if bytes.HasSuffix(w.ansiseq.Bytes(), []byte("[0m")) {
  29. // reset sequence
  30. w.lastseq.Reset()
  31. w.seqchanged = false
  32. } else if c == 'm' {
  33. // color code
  34. _, _ = w.lastseq.Write(w.ansiseq.Bytes())
  35. }
  36. _, _ = w.ansiseq.WriteTo(w.Forward)
  37. }
  38. } else {
  39. _, err := w.writeRune(c)
  40. if err != nil {
  41. return 0, err
  42. }
  43. }
  44. }
  45. return len(b), nil
  46. }
  47. func (w *Writer) writeRune(r rune) (int, error) {
  48. if w.runeBuf == nil {
  49. w.runeBuf = make([]byte, utf8.UTFMax)
  50. }
  51. n := utf8.EncodeRune(w.runeBuf, r)
  52. return w.Forward.Write(w.runeBuf[:n])
  53. }
  54. func (w *Writer) LastSequence() string {
  55. return w.lastseq.String()
  56. }
  57. func (w *Writer) ResetAnsi() {
  58. if !w.seqchanged {
  59. return
  60. }
  61. _, _ = w.Forward.Write([]byte("\x1b[0m"))
  62. }
  63. func (w *Writer) RestoreAnsi() {
  64. _, _ = w.Forward.Write(w.lastseq.Bytes())
  65. }