wordwrap.go 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167
  1. package wordwrap
  2. import (
  3. "bytes"
  4. "strings"
  5. "unicode"
  6. "github.com/muesli/reflow/ansi"
  7. )
  8. var (
  9. defaultBreakpoints = []rune{'-'}
  10. defaultNewline = []rune{'\n'}
  11. )
  12. // WordWrap contains settings and state for customisable text reflowing with
  13. // support for ANSI escape sequences. This means you can style your terminal
  14. // output without affecting the word wrapping algorithm.
  15. type WordWrap struct {
  16. Limit int
  17. Breakpoints []rune
  18. Newline []rune
  19. KeepNewlines bool
  20. buf bytes.Buffer
  21. space bytes.Buffer
  22. word ansi.Buffer
  23. lineLen int
  24. ansi bool
  25. }
  26. // NewWriter returns a new instance of a word-wrapping writer, initialized with
  27. // default settings.
  28. func NewWriter(limit int) *WordWrap {
  29. return &WordWrap{
  30. Limit: limit,
  31. Breakpoints: defaultBreakpoints,
  32. Newline: defaultNewline,
  33. KeepNewlines: true,
  34. }
  35. }
  36. // Bytes is shorthand for declaring a new default WordWrap instance,
  37. // used to immediately word-wrap a byte slice.
  38. func Bytes(b []byte, limit int) []byte {
  39. f := NewWriter(limit)
  40. _, _ = f.Write(b)
  41. _ = f.Close()
  42. return f.Bytes()
  43. }
  44. // String is shorthand for declaring a new default WordWrap instance,
  45. // used to immediately word-wrap a string.
  46. func String(s string, limit int) string {
  47. return string(Bytes([]byte(s), limit))
  48. }
  49. func (w *WordWrap) addSpace() {
  50. w.lineLen += w.space.Len()
  51. _, _ = w.buf.Write(w.space.Bytes())
  52. w.space.Reset()
  53. }
  54. func (w *WordWrap) addWord() {
  55. if w.word.Len() > 0 {
  56. w.addSpace()
  57. w.lineLen += w.word.PrintableRuneWidth()
  58. _, _ = w.buf.Write(w.word.Bytes())
  59. w.word.Reset()
  60. }
  61. }
  62. func (w *WordWrap) addNewLine() {
  63. _, _ = w.buf.WriteRune('\n')
  64. w.lineLen = 0
  65. w.space.Reset()
  66. }
  67. func inGroup(a []rune, c rune) bool {
  68. for _, v := range a {
  69. if v == c {
  70. return true
  71. }
  72. }
  73. return false
  74. }
  75. // Write is used to write more content to the word-wrap buffer.
  76. func (w *WordWrap) Write(b []byte) (int, error) {
  77. if w.Limit == 0 {
  78. return w.buf.Write(b)
  79. }
  80. s := string(b)
  81. if !w.KeepNewlines {
  82. s = strings.Replace(strings.TrimSpace(s), "\n", " ", -1)
  83. }
  84. for _, c := range s {
  85. if c == '\x1B' {
  86. // ANSI escape sequence
  87. _, _ = w.word.WriteRune(c)
  88. w.ansi = true
  89. } else if w.ansi {
  90. _, _ = w.word.WriteRune(c)
  91. if (c >= 0x40 && c <= 0x5a) || (c >= 0x61 && c <= 0x7a) {
  92. // ANSI sequence terminated
  93. w.ansi = false
  94. }
  95. } else if inGroup(w.Newline, c) {
  96. // end of current line
  97. // see if we can add the content of the space buffer to the current line
  98. if w.word.Len() == 0 {
  99. if w.lineLen+w.space.Len() > w.Limit {
  100. w.lineLen = 0
  101. } else {
  102. // preserve whitespace
  103. _, _ = w.buf.Write(w.space.Bytes())
  104. }
  105. w.space.Reset()
  106. }
  107. w.addWord()
  108. w.addNewLine()
  109. } else if unicode.IsSpace(c) {
  110. // end of current word
  111. w.addWord()
  112. _, _ = w.space.WriteRune(c)
  113. } else if inGroup(w.Breakpoints, c) {
  114. // valid breakpoint
  115. w.addSpace()
  116. w.addWord()
  117. _, _ = w.buf.WriteRune(c)
  118. } else {
  119. // any other character
  120. _, _ = w.word.WriteRune(c)
  121. // add a line break if the current word would exceed the line's
  122. // character limit
  123. if w.lineLen+w.space.Len()+w.word.PrintableRuneWidth() > w.Limit &&
  124. w.word.PrintableRuneWidth() < w.Limit {
  125. w.addNewLine()
  126. }
  127. }
  128. }
  129. return len(b), nil
  130. }
  131. // Close will finish the word-wrap operation. Always call it before trying to
  132. // retrieve the final result.
  133. func (w *WordWrap) Close() error {
  134. w.addWord()
  135. return nil
  136. }
  137. // Bytes returns the word-wrapped result as a byte slice.
  138. func (w *WordWrap) Bytes() []byte {
  139. return w.buf.Bytes()
  140. }
  141. // String returns the word-wrapped result as a string.
  142. func (w *WordWrap) String() string {
  143. return w.buf.String()
  144. }