1package tree 2 3import ( 4 "bytes" 5) 6 7// Node is a tree node 8type Node struct { 9 Name string 10 Type string 11 Template bool 12 Mount bool 13 Path string 14 Subtree *Tree 15} 16 17const ( 18 // INF allows to have a full recursion until the leaves of a tree 19 INF = -1 20) 21 22// Nodes is a slice of nodes which can be sorted 23type Nodes []*Node 24 25func (n Nodes) Len() int { 26 return len(n) 27} 28 29func (n Nodes) Less(i, j int) bool { 30 return n[i].Name < n[j].Name 31} 32 33func (n Nodes) Swap(i, j int) { 34 n[i], n[j] = n[j], n[i] 35} 36 37// Equals compares to another node 38func (n Node) Equals(other Node) bool { 39 if n.Name != other.Name { 40 return false 41 } 42 if n.Type != other.Type { 43 return false 44 } 45 if n.Subtree != nil { 46 if other.Subtree == nil { 47 return false 48 } 49 if !n.Subtree.Equals(other.Subtree) { 50 return false 51 } 52 } else { 53 if other.Subtree != nil { 54 return false 55 } 56 } 57 return true 58} 59 60// format returns a pretty printed string of all nodes in and below 61// this node, e.g. ├── baz 62func (n *Node) format(prefix string, last bool, maxDepth, curDepth int) string { 63 if maxDepth > INF && (curDepth > maxDepth+1) { 64 return "" 65 } 66 67 out := bytes.NewBufferString(prefix) 68 // adding either an L or a T, depending if this is the last node 69 // or not 70 if last { 71 _, _ = out.WriteString(symLeaf) 72 } else { 73 _, _ = out.WriteString(symBranch) 74 } 75 // the next levels prefix needs to be extended depending if 76 // this is the last node in a group or not 77 if last { 78 prefix += symEmpty 79 } else { 80 prefix += symVert 81 } 82 83 // any mount will be colored and include the on-disk path 84 switch { 85 case n.Mount: 86 _, _ = out.WriteString(colMount(n.Name + " (" + n.Path + ")")) 87 case n.Type == "dir": 88 _, _ = out.WriteString(colDir(n.Name + sep)) 89 default: 90 _, _ = out.WriteString(n.Name) 91 } 92 // mark templates 93 if n.Template { 94 _, _ = out.WriteString(" " + colTpl("(template)")) 95 } 96 // finish this output 97 _, _ = out.WriteString("\n") 98 99 if n.Subtree == nil { 100 return out.String() 101 } 102 103 // let our children format themselves 104 for i, node := range n.Subtree.Nodes { 105 last := i == len(n.Subtree.Nodes)-1 106 _, _ = out.WriteString(node.format(prefix, last, maxDepth, curDepth+1)) 107 } 108 return out.String() 109} 110 111// Len returns the length of this subtree 112func (n *Node) Len() int { 113 if n.Type == "file" { 114 return 1 115 } 116 var l int 117 for _, t := range n.Subtree.Nodes { 118 l += t.Len() 119 } 120 return l 121} 122 123func (n *Node) list(prefix string, maxDepth, curDepth int, files bool) []string { 124 if maxDepth >= 0 && curDepth > maxDepth { 125 return nil 126 } 127 128 if prefix != "" { 129 prefix += sep 130 } 131 prefix += n.Name 132 133 // if it's a file and we are looking for files 134 if n.Type == "file" && files { 135 // we return the file 136 return []string{prefix} 137 } else if curDepth == maxDepth && n.Type != "file" { 138 // otherwise if we are "at the bottom" and it's not a file 139 // we return the directory name with a separator at the end 140 return []string{prefix + sep} 141 } 142 143 out := make([]string, 0, n.Len()) 144 // if we don't have subitems, then it's a leaf and we return 145 // (notice that this is what ends the recursion when maxDepth is set to -1) 146 if n.Subtree == nil { 147 return out 148 } 149 150 // this is the part that will list the subdirectories on their own line when using the -d option 151 if !files { 152 out = append(out, prefix+sep) 153 } 154 155 // we keep listing the subtree nodes if we haven't exited yet. 156 for _, t := range n.Subtree.Nodes { 157 out = append(out, t.list(prefix, maxDepth, curDepth+1, files)...) 158 } 159 return out 160} 161