print.go 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304
  1. package pretty
  2. import (
  3. "bytes"
  4. "fmt"
  5. "io"
  6. "github.com/fatih/color"
  7. )
  8. // NB(directxman12): this isn't particularly elegant, but it's also
  9. // sufficiently simple as to be maintained here. Man (roff) would've
  10. // probably worked, but it's not necessarily on Windows by default.
  11. // Span is a chunk of content that is writable to an output, but knows how to
  12. // calculate its apparent visual "width" on the terminal (not to be confused
  13. // with the raw length, which may include zero-width coloring sequences).
  14. type Span interface {
  15. // VisualLength reports the "width" as perceived by the user on the terminal
  16. // (i.e. widest line, ignoring ANSI escape characters).
  17. VisualLength() int
  18. // WriteTo writes the full span contents to the given writer.
  19. WriteTo(io.Writer) error
  20. }
  21. // Table is a Span that writes its data in table form, with sizing controlled
  22. // by the given table calculator. Rows are started with StartRow, followed by
  23. // some calls to Column, followed by a call to EndRow. Once all rows are
  24. // added, the table can be used as a Span.
  25. type Table struct {
  26. Sizing *TableCalculator
  27. cellsByRow [][]Span
  28. colSizes []int
  29. }
  30. // StartRow starts a new row.
  31. // It must eventually be followed by EndRow.
  32. func (t *Table) StartRow() {
  33. t.cellsByRow = append(t.cellsByRow, []Span(nil))
  34. }
  35. // EndRow ends the currently started row.
  36. func (t *Table) EndRow() {
  37. lastRow := t.cellsByRow[len(t.cellsByRow)-1]
  38. sizes := make([]int, len(lastRow))
  39. for i, cell := range lastRow {
  40. sizes[i] = cell.VisualLength()
  41. }
  42. t.Sizing.AddRowSizes(sizes...)
  43. }
  44. // Column adds the given span as a new column to the current row.
  45. func (t *Table) Column(contents Span) {
  46. currentRowInd := len(t.cellsByRow) - 1
  47. t.cellsByRow[currentRowInd] = append(t.cellsByRow[currentRowInd], contents)
  48. }
  49. // SkipRow prints a span without having it contribute to the table calculation.
  50. func (t *Table) SkipRow(contents Span) {
  51. t.cellsByRow = append(t.cellsByRow, []Span{contents})
  52. }
  53. func (t *Table) WriteTo(out io.Writer) error {
  54. if t.colSizes == nil {
  55. t.colSizes = t.Sizing.ColumnWidths()
  56. }
  57. for _, cells := range t.cellsByRow {
  58. currentPosition := 0
  59. for colInd, cell := range cells {
  60. colSize := t.colSizes[colInd]
  61. diff := colSize - cell.VisualLength()
  62. if err := cell.WriteTo(out); err != nil {
  63. return err
  64. }
  65. if diff > 0 {
  66. if err := writePadding(out, columnPadding, diff); err != nil {
  67. return err
  68. }
  69. }
  70. currentPosition += colSize
  71. }
  72. if _, err := fmt.Fprint(out, "\n"); err != nil {
  73. return err
  74. }
  75. }
  76. return nil
  77. }
  78. func (t *Table) VisualLength() int {
  79. if t.colSizes == nil {
  80. t.colSizes = t.Sizing.ColumnWidths()
  81. }
  82. res := 0
  83. for _, colSize := range t.colSizes {
  84. res += colSize
  85. }
  86. return res
  87. }
  88. // Text is a span that simply contains raw text. It's a good starting point.
  89. type Text string
  90. func (t Text) VisualLength() int { return len(t) }
  91. func (t Text) WriteTo(w io.Writer) error {
  92. _, err := w.Write([]byte(t))
  93. return err
  94. }
  95. // indented is a span that indents all lines by the given number of tabs.
  96. type indented struct {
  97. Amount int
  98. Content Span
  99. }
  100. func (i *indented) VisualLength() int { return i.Content.VisualLength() }
  101. func (i *indented) WriteTo(w io.Writer) error {
  102. var out bytes.Buffer
  103. if err := i.Content.WriteTo(&out); err != nil {
  104. return err
  105. }
  106. lines := bytes.Split(out.Bytes(), []byte("\n"))
  107. for lineInd, line := range lines {
  108. if lineInd != 0 {
  109. if _, err := w.Write([]byte("\n")); err != nil {
  110. return err
  111. }
  112. }
  113. if len(line) == 0 {
  114. continue
  115. }
  116. if err := writePadding(w, indentPadding, i.Amount); err != nil {
  117. return err
  118. }
  119. if _, err := w.Write(line); err != nil {
  120. return err
  121. }
  122. }
  123. return nil
  124. }
  125. // Indented returns a span that indents all lines by the given number of tabs.
  126. func Indented(amt int, content Span) Span {
  127. return &indented{Amount: amt, Content: content}
  128. }
  129. // fromWriter is a span that takes content from a function expecting a Writer.
  130. type fromWriter struct {
  131. cache []byte
  132. cacheError error
  133. run func(io.Writer) error
  134. }
  135. func (f *fromWriter) VisualLength() int {
  136. if f.cache == nil {
  137. var buf bytes.Buffer
  138. if err := f.run(&buf); err != nil {
  139. f.cacheError = err
  140. }
  141. f.cache = buf.Bytes()
  142. }
  143. return len(f.cache)
  144. }
  145. func (f *fromWriter) WriteTo(w io.Writer) error {
  146. if f.cache != nil {
  147. if f.cacheError != nil {
  148. return f.cacheError
  149. }
  150. _, err := w.Write(f.cache)
  151. return err
  152. }
  153. return f.run(w)
  154. }
  155. // FromWriter returns a span that takes content from a function expecting a Writer.
  156. func FromWriter(run func(io.Writer) error) Span {
  157. return &fromWriter{run: run}
  158. }
  159. // Decoration represents a terminal decoration.
  160. type Decoration color.Color
  161. // Containing returns a Span that has the given decoration applied.
  162. func (d Decoration) Containing(contents Span) Span {
  163. return &decorated{
  164. Contents: contents,
  165. Attributes: color.Color(d),
  166. }
  167. }
  168. // decorated is a span that has some terminal decoration applied.
  169. type decorated struct {
  170. Contents Span
  171. Attributes color.Color
  172. }
  173. func (d *decorated) VisualLength() int { return d.Contents.VisualLength() }
  174. func (d *decorated) WriteTo(w io.Writer) error {
  175. oldOut := color.Output
  176. color.Output = w
  177. defer func() { color.Output = oldOut }()
  178. d.Attributes.Set()
  179. defer color.Unset()
  180. return d.Contents.WriteTo(w)
  181. }
  182. // SpanWriter is a span that contains multiple sub-spans.
  183. type SpanWriter struct {
  184. contents []Span
  185. }
  186. func (m *SpanWriter) VisualLength() int {
  187. res := 0
  188. for _, span := range m.contents {
  189. res += span.VisualLength()
  190. }
  191. return res
  192. }
  193. func (m *SpanWriter) WriteTo(w io.Writer) error {
  194. for _, span := range m.contents {
  195. if err := span.WriteTo(w); err != nil {
  196. return err
  197. }
  198. }
  199. return nil
  200. }
  201. // Print adds a new span to this SpanWriter.
  202. func (m *SpanWriter) Print(s Span) {
  203. m.contents = append(m.contents, s)
  204. }
  205. // lines is a span that adds some newlines, optionally followed by some content.
  206. type lines struct {
  207. content Span
  208. amountBefore int
  209. }
  210. func (l *lines) VisualLength() int {
  211. if l.content == nil {
  212. return 0
  213. }
  214. return l.content.VisualLength()
  215. }
  216. func (l *lines) WriteTo(w io.Writer) error {
  217. if err := writePadding(w, linesPadding, l.amountBefore); err != nil {
  218. return err
  219. }
  220. if l.content != nil {
  221. if err := l.content.WriteTo(w); err != nil {
  222. return err
  223. }
  224. }
  225. return nil
  226. }
  227. // Newlines returns a span just containing some newlines.
  228. func Newlines(amt int) Span {
  229. return &lines{amountBefore: amt}
  230. }
  231. // Line returns a span that emits a newline, followed by the given content.
  232. func Line(content Span) Span {
  233. return &lines{amountBefore: 1, content: content}
  234. }
  235. var (
  236. columnPadding = []byte(" ")
  237. indentPadding = []byte("\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t")
  238. linesPadding = []byte("\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n")
  239. )
  240. // writePadding writes out padding of the given type in the given amount to the writer.
  241. // Each byte in the padding buffer contributes 1 to the amount -- the padding being
  242. // a buffer is just for efficiency.
  243. func writePadding(out io.Writer, typ []byte, amt int) error {
  244. if amt <= len(typ) {
  245. _, err := out.Write(typ[:amt])
  246. return err
  247. }
  248. num := amt / len(typ)
  249. rem := amt % len(typ)
  250. for i := 0; i < num; i++ {
  251. if _, err := out.Write(typ); err != nil {
  252. return err
  253. }
  254. }
  255. if _, err := out.Write(typ[:rem]); err != nil {
  256. return err
  257. }
  258. return nil
  259. }