indent.go 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. package text
  2. import (
  3. "io"
  4. )
  5. // Indent inserts prefix at the beginning of each non-empty line of s. The
  6. // end-of-line marker is NL.
  7. func Indent(s, prefix string) string {
  8. return string(IndentBytes([]byte(s), []byte(prefix)))
  9. }
  10. // IndentBytes inserts prefix at the beginning of each non-empty line of b.
  11. // The end-of-line marker is NL.
  12. func IndentBytes(b, prefix []byte) []byte {
  13. var res []byte
  14. bol := true
  15. for _, c := range b {
  16. if bol && c != '\n' {
  17. res = append(res, prefix...)
  18. }
  19. res = append(res, c)
  20. bol = c == '\n'
  21. }
  22. return res
  23. }
  24. // Writer indents each line of its input.
  25. type indentWriter struct {
  26. w io.Writer
  27. bol bool
  28. pre [][]byte
  29. sel int
  30. off int
  31. }
  32. // NewIndentWriter makes a new write filter that indents the input
  33. // lines. Each line is prefixed in order with the corresponding
  34. // element of pre. If there are more lines than elements, the last
  35. // element of pre is repeated for each subsequent line.
  36. func NewIndentWriter(w io.Writer, pre ...[]byte) io.Writer {
  37. return &indentWriter{
  38. w: w,
  39. pre: pre,
  40. bol: true,
  41. }
  42. }
  43. // The only errors returned are from the underlying indentWriter.
  44. func (w *indentWriter) Write(p []byte) (n int, err error) {
  45. for _, c := range p {
  46. if w.bol {
  47. var i int
  48. i, err = w.w.Write(w.pre[w.sel][w.off:])
  49. w.off += i
  50. if err != nil {
  51. return n, err
  52. }
  53. }
  54. _, err = w.w.Write([]byte{c})
  55. if err != nil {
  56. return n, err
  57. }
  58. n++
  59. w.bol = c == '\n'
  60. if w.bol {
  61. w.off = 0
  62. if w.sel < len(w.pre)-1 {
  63. w.sel++
  64. }
  65. }
  66. }
  67. return n, nil
  68. }