1package tree 2 3import ( 4 "bytes" 5 "fmt" 6 "path/filepath" 7 "strings" 8 9 "github.com/fatih/color" 10) 11 12const ( 13 symEmpty = " " 14 symBranch = "├── " 15 symLeaf = "└── " 16 symVert = "│ " 17) 18 19var ( 20 colMount = color.New(color.FgCyan, color.Bold).SprintfFunc() 21 colDir = color.New(color.FgBlue, color.Bold).SprintfFunc() 22 colTpl = color.New(color.FgGreen, color.Bold).SprintfFunc() 23 // sep is intentionally NOT platform-agnostic. This is used for the CLI output 24 // and should always be a regular slash. 25 sep = "/" 26) 27 28// Root is the root of a tree 29type Root struct { 30 Name string 31 Subtree *Tree 32 Prefix string 33} 34 35// New creates a new tree 36func New(name string) *Root { 37 return &Root{ 38 Name: name, 39 Subtree: NewTree(), 40 } 41} 42 43// AddFile adds a new file to the tree 44func (r *Root) AddFile(path string, _ string) error { 45 return r.insert(path, false, "") 46} 47 48// AddMount adds a new mount point to the tree 49func (r *Root) AddMount(path, dest string) error { 50 return r.insert(path, false, dest) 51} 52 53// AddTemplate adds a template to the tree 54func (r *Root) AddTemplate(path string) error { 55 return r.insert(path, true, "") 56} 57 58func (r *Root) insert(path string, template bool, mountPath string) error { 59 t := r.Subtree 60 p := strings.Split(path, "/") 61 for i, e := range p { 62 n := &Node{ 63 Name: e, 64 Type: "dir", 65 Subtree: NewTree(), 66 } 67 if i == len(p)-1 { 68 n.Type = "file" 69 n.Subtree = nil 70 n.Template = template 71 if mountPath != "" { 72 n.Mount = true 73 n.Path = mountPath 74 } 75 } 76 node, _ := t.Insert(n) 77 // do we need to extend an existing subtree? 78 if i < len(p)-1 && node.Subtree == nil { 79 node.Subtree = NewTree() 80 node.Type = "dir" 81 } 82 t = node.Subtree 83 } 84 return nil 85} 86 87// Format returns a pretty printed string of all nodes in and below 88// this node, e.g. ├── baz 89func (r *Root) Format(maxDepth int) string { 90 out := &bytes.Buffer{} 91 92 // any mount will be colored and include the on-disk path 93 _, _ = out.WriteString(colDir(r.Name)) 94 95 // finish this folders output 96 _, _ = out.WriteString("\n") 97 98 // let our children format themselves 99 for i, node := range r.Subtree.Nodes { 100 last := i == len(r.Subtree.Nodes)-1 101 _, _ = out.WriteString(node.format("", last, maxDepth, 1)) 102 } 103 return out.String() 104} 105 106// List returns a flat list of all files in this tree 107func (r *Root) List(maxDepth int) []string { 108 out := make([]string, 0, r.Len()) 109 for _, t := range r.Subtree.Nodes { 110 out = append(out, t.list(r.Prefix, maxDepth, 0, true)...) 111 } 112 return out 113} 114 115// ListFolders returns a flat list of all folders in this tree 116func (r *Root) ListFolders(maxDepth int) []string { 117 out := make([]string, 0, r.Len()) 118 for _, t := range r.Subtree.Nodes { 119 out = append(out, t.list(r.Prefix, maxDepth, 0, false)...) 120 } 121 return out 122} 123 124// String returns the name of this tree 125func (r *Root) String() string { 126 return r.Name 127} 128 129// FindFolder returns the subtree rooted at path 130func (r *Root) FindFolder(path string) (*Root, error) { 131 path = strings.TrimSuffix(path, "/") 132 t := r.Subtree 133 p := strings.Split(path, "/") 134 prefix := "" 135 for _, e := range p { 136 _, node := t.find(e) 137 if node == nil || node.Type == "file" || node.Subtree == nil { 138 return nil, fmt.Errorf("not found") 139 } 140 t = node.Subtree 141 prefix = filepath.Join(prefix, e) 142 } 143 return &Root{Name: r.Name, Subtree: t, Prefix: prefix}, nil 144} 145 146// SetName changes the name of this tree 147func (r *Root) SetName(n string) { 148 r.Name = n 149} 150 151// Len returns the number of entries in this folder and all subfolder including 152// this folder itself 153func (r *Root) Len() int { 154 var l int 155 for _, t := range r.Subtree.Nodes { 156 l += t.Len() 157 } 158 return l 159} 160