1 package drawing
2
3 import (
4 "bytes"
5 "fmt"
6 "io"
7 "os"
8 "sort"
9 "strings"
10 )
11
12 const (
13 lightVerticalRight = "\u251C"
14 lightUpRight = "\u2514"
15 lightHorizontal = "\u2500"
16 lightVertical = "\u2502"
17 )
18
19 type StringTree struct {
20 Data string
21 Children []*StringTree
22 Labels map[string]string
23 }
24
25 func (t *StringTree) String() string {
26 buf := &bytes.Buffer{}
27 t.Fprint(buf)
28 return buf.String()
29 }
30
31 func (t *StringTree) Print() {
32 t.Fprint(os.Stdout)
33 }
34
35 func (t *StringTree) Fprint(w io.Writer) {
36 fmt.Fprintln(w, t.Data, t.labels())
37 t.printTree(w, -1, 0, 0, "", true)
38 }
39
40 func (t *StringTree) printTree(w io.Writer, pi, pn, depth int, indention string, sorted bool) {
41 if sorted {
42 sort.Slice(t.Children, func(i int, j int) bool {
43 return t.Children[i].depth(0) < t.Children[j].depth(0)
44 })
45 }
46
47 n := len(t.Children)
48 for i, tree := range t.Children {
49 var marker string
50 if i == len(t.Children)-1 {
51 marker = lightUpRight + lightHorizontal
52 } else {
53 marker = lightVerticalRight + lightHorizontal
54 }
55 debugStr := ""
56 debug := false
57 if debug {
58 maxDepth := tree.depth(0)
59 debugStr = fmt.Sprintf(" [depth: %d, indent: '%s', i: %d, n: %d, pi: %d, pn: %d, maxDepth: %d]", depth, indention, i, n, pi, pn, maxDepth)
60 }
61
62 fmt.Fprintf(w, "%s%s %s%s %v\n", indention, marker, tree.Data, debugStr, tree.labels())
63 tree.printTree(w, i, n, depth+1, indention+indent(i, n), sorted)
64 }
65 }
66
67 func (t *StringTree) depth(depth int) int {
68 max := depth
69 for _, tree := range t.Children {
70 treeD := tree.depth(depth + 1)
71 if treeD > max {
72 max = treeD
73 }
74 }
75 return max
76 }
77
78 func indent(i, n int) string {
79 s := ""
80 if i != n-1 {
81 s += lightVertical + " "
82 } else {
83 s += " "
84 }
85 return s
86 }
87
88 func (t *StringTree) labels() string {
89 labelStrs := []string{}
90 for k, v := range t.Labels {
91 if v == "" {
92 labelStrs = append(labelStrs, fmt.Sprintf("%v", k))
93 } else {
94 labelStrs = append(labelStrs, fmt.Sprintf("%v=%v", k, v))
95 }
96 }
97 if len(labelStrs) > 0 {
98 return "[" + strings.Join(labelStrs, ", ") + "]"
99 }
100 return ""
101 }
102
View as plain text