Line data Source code
1 : // Copyright 2017 The Cockroach Authors.
2 : //
3 : // Use of this software is governed by the Business Source License
4 : // included in the file licenses/BSL.txt.
5 : //
6 : // As of the Change Date specified in that file, in accordance with
7 : // the Business Source License, use of this software will be governed
8 : // by the Apache License, Version 2.0, included in the file
9 : // licenses/APL.txt.
10 :
11 : package treeprinter
12 :
13 : import (
14 : "bytes"
15 : "fmt"
16 : "strings"
17 : )
18 :
19 : var (
20 : edgeLinkChr = rune('│')
21 : edgeMidChr = rune('├')
22 : edgeLastChr = rune('└')
23 : horLineChr = rune('─')
24 : bulletChr = rune('•')
25 : )
26 :
27 : // Node is a handle associated with a specific depth in a tree. See below for
28 : // sample usage.
29 : type Node struct {
30 : tree *tree
31 : level int
32 : }
33 :
34 : // New creates a tree printer and returns a sentinel node reference which
35 : // should be used to add the root. Sample usage:
36 : //
37 : // tp := New()
38 : // root := tp.Child("root")
39 : // root.Child("child-1")
40 : // root.Child("child-2").Child("grandchild\ngrandchild-more-info")
41 : // root.Child("child-3")
42 : //
43 : // fmt.Print(tp.String())
44 : //
45 : // Output:
46 : //
47 : // root
48 : // ├── child-1
49 : // ├── child-2
50 : // │ └── grandchild
51 : // │ grandchild-more-info
52 : // └── child-3
53 : //
54 : // Note that the Child calls can't be rearranged arbitrarily; they have
55 : // to be in the order they need to be displayed (depth-first pre-order).
56 1 : func New() Node {
57 1 : return NewWithStyle(DefaultStyle)
58 1 : }
59 :
60 : // NewWithStyle creates a tree printer like New, permitting customization of
61 : // the style of the resulting tree.
62 1 : func NewWithStyle(style Style) Node {
63 1 : t := &tree{style: style}
64 1 :
65 1 : switch style {
66 1 : case CompactStyle:
67 1 : t.edgeLink = []rune{edgeLinkChr}
68 1 : t.edgeMid = []rune{edgeMidChr, ' '}
69 1 : t.edgeLast = []rune{edgeLastChr, ' '}
70 :
71 1 : case BulletStyle:
72 1 : t.edgeLink = []rune{edgeLinkChr}
73 1 : t.edgeMid = []rune{edgeMidChr, horLineChr, horLineChr, ' '}
74 1 : t.edgeLast = []rune{edgeLastChr, horLineChr, horLineChr, ' '}
75 :
76 1 : default:
77 1 : t.edgeLink = []rune{' ', edgeLinkChr}
78 1 : t.edgeMid = []rune{' ', edgeMidChr, horLineChr, horLineChr, ' '}
79 1 : t.edgeLast = []rune{' ', edgeLastChr, horLineChr, horLineChr, ' '}
80 : }
81 :
82 1 : return Node{
83 1 : tree: t,
84 1 : level: 0,
85 1 : }
86 : }
87 :
88 : // Style is one of the predefined treeprinter styles.
89 : type Style int
90 :
91 : const (
92 : // DefaultStyle is the default style. Example:
93 : //
94 : // foo
95 : // ├── bar1
96 : // │ bar2
97 : // │ └── baz
98 : // └── qux
99 : //
100 : DefaultStyle Style = iota
101 :
102 : // CompactStyle is a compact style, for deeper trees. Example:
103 : //
104 : // foo
105 : // ├ bar1
106 : // │ bar2
107 : // │ └ baz
108 : // └ qux
109 : //
110 : CompactStyle
111 :
112 : // BulletStyle is a style that shows a bullet for each node, and groups any
113 : // other lines under that bullet. Example:
114 : //
115 : // • foo
116 : // │
117 : // ├── • bar1
118 : // │ │ bar2
119 : // │ │
120 : // │ └── • baz
121 : // │
122 : // └── • qux
123 : //
124 : BulletStyle
125 : )
126 :
127 : // tree implements the tree printing machinery.
128 : //
129 : // All Nodes hold a reference to the tree and Node calls result in modification
130 : // of the tree. At any point in time, tree.rows contains the formatted tree that
131 : // was described by the Node calls performed so far.
132 : //
133 : // When new nodes are added, some of the characters of the previous formatted
134 : // tree need to be updated. Here is an example stepping through the state:
135 : //
136 : // API call Rows
137 : //
138 : //
139 : // tp := New() <empty>
140 : //
141 : //
142 : // root := tp.Child("root") root
143 : //
144 : //
145 : // root.Child("child-1") root
146 : // └── child-1
147 : //
148 : //
149 : // c2 := root.Child("child-2") root
150 : // ├── child-1
151 : // └── child-2
152 : //
153 : // Note: here we had to go back up and change └─ into ├─ for child-1.
154 : //
155 : //
156 : // c2.Child("grandchild") root
157 : // ├── child-1
158 : // └── child-2
159 : // └── grandchild
160 : //
161 : //
162 : // root.Child("child-3" root
163 : // ├── child-1
164 : // ├── child-2
165 : // │ └── grandchild
166 : // └── child-3
167 : //
168 : // Note: here we had to go back up and change └─ into ├─ for child-2, and
169 : // add a │ on the grandchild row. In general, we may need to add an
170 : // arbitrary number of vertical bars.
171 : //
172 : // In order to perform these character changes, we maintain information about
173 : // the nodes on the bottom-most path.
174 : type tree struct {
175 : style Style
176 :
177 : // rows maintains the rows accumulated so far, as rune arrays.
178 : rows [][]rune
179 :
180 : // stack contains information pertaining to the nodes on the bottom-most path
181 : // of the tree.
182 : stack []nodeInfo
183 :
184 : edgeLink []rune
185 : edgeMid []rune
186 : edgeLast []rune
187 : }
188 :
189 : type nodeInfo struct {
190 : // firstChildConnectRow is the index (in tree.rows) of the row up to which we
191 : // have to connect the first child of this node.
192 : firstChildConnectRow int
193 :
194 : // nextSiblingConnectRow is the index (in tree.rows) of the row up to which we
195 : // have to connect the next sibling of this node. Typically this is the same
196 : // with firstChildConnectRow, except when the node has multiple rows. For
197 : // example:
198 : //
199 : // foo
200 : // └── bar1 <---- nextSiblingConnectRow
201 : // bar2 <---- firstChildConnectRow
202 : //
203 : // firstChildConnectRow is used when adding "baz", nextSiblingConnectRow
204 : // is used when adding "qux":
205 : // foo
206 : // ├── bar1
207 : // │ bar2
208 : // │ └── baz
209 : // └── qux
210 : //
211 : nextSiblingConnectRow int
212 : }
213 :
214 : // set copies the string of runes into a given row, at a specific position. The
215 : // row is extended with spaces if needed.
216 1 : func (t *tree) set(rowIdx int, colIdx int, what []rune) {
217 1 : // Extend the line if necessary.
218 1 : for len(t.rows[rowIdx]) < colIdx+len(what) {
219 1 : t.rows[rowIdx] = append(t.rows[rowIdx], ' ')
220 1 : }
221 1 : copy(t.rows[rowIdx][colIdx:], what)
222 : }
223 :
224 : // addRow adds a row with a given text, with the proper indentation for the
225 : // given level.
226 1 : func (t *tree) addRow(level int, text string) (rowIdx int) {
227 1 : runes := []rune(text)
228 1 : // Each level indents by this much.
229 1 : k := len(t.edgeLast)
230 1 : indent := level * k
231 1 : row := make([]rune, indent+len(runes))
232 1 : for i := 0; i < indent; i++ {
233 1 : row[i] = ' '
234 1 : }
235 1 : copy(row[indent:], runes)
236 1 : t.rows = append(t.rows, row)
237 1 : return len(t.rows) - 1
238 : }
239 :
240 : // Childf adds a node as a child of the given node.
241 1 : func (n Node) Childf(format string, args ...interface{}) Node {
242 1 : return n.Child(fmt.Sprintf(format, args...))
243 1 : }
244 :
245 : // Child adds a node as a child of the given node. Multi-line strings are
246 : // supported with appropriate indentation.
247 1 : func (n Node) Child(text string) Node {
248 1 : if strings.ContainsRune(text, '\n') {
249 1 : splitLines := strings.Split(text, "\n")
250 1 : node := n.childLine(splitLines[0])
251 1 : for _, l := range splitLines[1:] {
252 1 : node.AddLine(l)
253 1 : }
254 1 : return node
255 : }
256 1 : return n.childLine(text)
257 : }
258 :
259 : // AddLine adds a new line to a node without an edge.
260 1 : func (n Node) AddLine(text string) {
261 1 : t := n.tree
262 1 : if t.style == BulletStyle {
263 1 : text = " " + text
264 1 : }
265 1 : rowIdx := t.addRow(n.level-1, text)
266 1 : if t.style != BulletStyle {
267 1 : t.stack[n.level-1].firstChildConnectRow = rowIdx
268 1 : }
269 : }
270 :
271 : // childLine adds a node as a child of the given node.
272 1 : func (n Node) childLine(text string) Node {
273 1 : t := n.tree
274 1 : if t.style == BulletStyle {
275 1 : text = fmt.Sprintf("%c %s", bulletChr, text)
276 1 : if n.level > 0 {
277 1 : n.AddEmptyLine()
278 1 : }
279 : }
280 1 : rowIdx := t.addRow(n.level, text)
281 1 : edgePos := (n.level - 1) * len(t.edgeLast)
282 1 : if n.level == 0 {
283 1 : // Case 1: root.
284 1 : if len(t.stack) != 0 {
285 0 : panic("multiple root nodes")
286 : }
287 1 : } else if len(t.stack) <= n.level {
288 1 : // Case 2: first child. Connect to parent.
289 1 : if len(t.stack) != n.level {
290 0 : panic("misuse of node")
291 : }
292 1 : parentRow := t.stack[n.level-1].firstChildConnectRow
293 1 : for i := parentRow + 1; i < rowIdx; i++ {
294 1 : t.set(i, edgePos, t.edgeLink)
295 1 : }
296 1 : t.set(rowIdx, edgePos, t.edgeLast)
297 1 : } else {
298 1 : // Case 3: non-first child. Connect to sibling.
299 1 : siblingRow := t.stack[n.level].nextSiblingConnectRow
300 1 : t.set(siblingRow, edgePos, t.edgeMid)
301 1 : for i := siblingRow + 1; i < rowIdx; i++ {
302 1 : t.set(i, edgePos, t.edgeLink)
303 1 : }
304 1 : t.set(rowIdx, edgePos, t.edgeLast)
305 1 : // Update the nextSiblingConnectRow.
306 1 : t.stack = t.stack[:n.level]
307 : }
308 :
309 1 : t.stack = append(t.stack, nodeInfo{
310 1 : firstChildConnectRow: rowIdx,
311 1 : nextSiblingConnectRow: rowIdx,
312 1 : })
313 1 :
314 1 : // Return a TreePrinter that can be used for children of this node.
315 1 : return Node{
316 1 : tree: t,
317 1 : level: n.level + 1,
318 1 : }
319 : }
320 :
321 : // AddEmptyLine adds an empty line to the output; used to introduce vertical
322 : // spacing as needed.
323 1 : func (n Node) AddEmptyLine() {
324 1 : n.tree.rows = append(n.tree.rows, []rune{})
325 1 : }
326 :
327 : // FormattedRows returns the formatted rows. Can only be called on the result of
328 : // treeprinter.New.
329 1 : func (n Node) FormattedRows() []string {
330 1 : if n.level != 0 {
331 0 : panic("Only the root can be stringified")
332 : }
333 1 : res := make([]string, len(n.tree.rows))
334 1 : for i, r := range n.tree.rows {
335 1 : res[i] = string(r)
336 1 : }
337 1 : return res
338 : }
339 :
340 1 : func (n Node) String() string {
341 1 : if n.level != 0 {
342 0 : panic("Only the root can be stringified")
343 : }
344 1 : var buf bytes.Buffer
345 1 : for _, r := range n.tree.rows {
346 1 : buf.WriteString(string(r))
347 1 : buf.WriteByte('\n')
348 1 : }
349 1 : return buf.String()
350 : }
|