1package tree 2 3import ( 4 "errors" 5 "fmt" 6 "io" 7 "os" 8 "os/user" 9 "path/filepath" 10 "regexp" 11 "sort" 12 "strconv" 13 "strings" 14) 15 16// Node represent some node in the tree 17// contains FileInfo, and its childs 18type Node struct { 19 os.FileInfo 20 path string 21 depth int 22 err error 23 nodes Nodes 24 vpaths map[string]bool 25} 26 27// List of nodes 28type Nodes []*Node 29 30// To use this package programmatically, you must implement this 31// interface. 32// For example: PTAL on 'cmd/tree/tree.go' 33type Fs interface { 34 Stat(path string) (os.FileInfo, error) 35 ReadDir(path string) ([]string, error) 36} 37 38// Options store the configuration for specific tree. 39// Note, that 'Fs', and 'OutFile' are required (OutFile can be os.Stdout). 40type Options struct { 41 Fs Fs 42 OutFile io.Writer 43 // List 44 All bool 45 DirsOnly bool 46 FullPath bool 47 IgnoreCase bool 48 FollowLink bool 49 DeepLevel int 50 Pattern string 51 IPattern string 52 MatchDirs bool 53 Prune bool 54 // File 55 ByteSize bool 56 UnitSize bool 57 FileMode bool 58 ShowUid bool 59 ShowGid bool 60 LastMod bool 61 Quotes bool 62 Inodes bool 63 Device bool 64 // Sort 65 NoSort bool 66 VerSort bool 67 ModSort bool 68 DirSort bool 69 NameSort bool 70 SizeSort bool 71 CTimeSort bool 72 ReverSort bool 73 // Graphics 74 NoIndent bool 75 Colorize bool 76 // Color defaults to ANSIColor() 77 Color func(*Node, string) string 78} 79 80func (opts *Options) color(node *Node, s string) string { 81 f := opts.Color 82 if f == nil { 83 f = ANSIColor 84 } 85 return f(node, s) 86} 87 88// New get path and create new node(root). 89func New(path string) *Node { 90 return &Node{path: path, vpaths: make(map[string]bool)} 91} 92 93// Visit all files under the given node. 94func (node *Node) Visit(opts *Options) (dirs, files int) { 95 // visited paths 96 if path, err := filepath.Abs(node.path); err == nil { 97 path = filepath.Clean(path) 98 node.vpaths[path] = true 99 } 100 // stat 101 fi, err := opts.Fs.Stat(node.path) 102 if err != nil { 103 node.err = err 104 return 105 } 106 node.FileInfo = fi 107 if !fi.IsDir() { 108 return 0, 1 109 } 110 // increase dirs only if it's a dir, but not the root. 111 if node.depth != 0 { 112 dirs++ 113 } 114 // DeepLevel option 115 if opts.DeepLevel > 0 && opts.DeepLevel <= node.depth { 116 return 117 } 118 // MatchDirs option 119 var dirMatch bool 120 if node.depth != 0 && opts.MatchDirs { 121 // then disable prune and pattern for immediate children 122 if opts.Pattern != "" { 123 dirMatch = node.match(opts.Pattern, opts) 124 } else if opts.IPattern != "" && node.match(opts.IPattern, opts) { 125 return 126 } 127 } 128 names, err := opts.Fs.ReadDir(node.path) 129 if err != nil { 130 node.err = err 131 return 132 } 133 node.nodes = make(Nodes, 0) 134 for _, name := range names { 135 // "all" option 136 if !opts.All && strings.HasPrefix(name, ".") { 137 continue 138 } 139 nnode := &Node{ 140 path: filepath.Join(node.path, name), 141 depth: node.depth + 1, 142 vpaths: node.vpaths, 143 } 144 d, f := nnode.Visit(opts) 145 if nnode.IsDir() { 146 // "prune" option, hide empty directories 147 if opts.Prune && f == 0 { 148 continue 149 } 150 if opts.MatchDirs && opts.IPattern != "" && nnode.match(opts.IPattern, opts) { 151 continue 152 } 153 } else if nnode.err == nil { 154 // "dirs only" option 155 if opts.DirsOnly { 156 continue 157 } 158 // Pattern matching 159 if !dirMatch && opts.Pattern != "" && !nnode.match(opts.Pattern, opts) { 160 continue 161 } 162 // IPattern matching 163 if opts.IPattern != "" && nnode.match(opts.IPattern, opts) { 164 continue 165 } 166 } 167 node.nodes = append(node.nodes, nnode) 168 dirs, files = dirs+d, files+f 169 } 170 // Sorting 171 if !opts.NoSort { 172 node.sort(opts) 173 } 174 return 175} 176 177func (node *Node) match(pattern string, opt *Options) bool { 178 var prefix string 179 if opt.IgnoreCase { 180 prefix = "(?i)" 181 } 182 search := node.Name() 183 if strings.Contains(pattern, "*") { 184 search = node.path 185 } 186 re, err := regexp.Compile(prefix + pattern) 187 return err == nil && re.FindString(search) != "" 188} 189 190func (node *Node) sort(opts *Options) { 191 var fn SortFunc 192 switch { 193 case opts.ModSort: 194 fn = ModSort 195 case opts.CTimeSort: 196 fn = CTimeSort 197 case opts.DirSort: 198 fn = DirSort 199 case opts.VerSort: 200 fn = VerSort 201 case opts.SizeSort: 202 fn = SizeSort 203 case opts.NameSort: 204 fn = NameSort 205 default: 206 fn = NameSort // Default should be sorted, not unsorted. 207 } 208 if fn != nil { 209 if opts.ReverSort { 210 sort.Sort(sort.Reverse(ByFunc{node.nodes, fn})) 211 } else { 212 sort.Sort(ByFunc{node.nodes, fn}) 213 } 214 } 215} 216 217// Path returns the Node's absolute path 218func (node *Node) Path() string { 219 return node.path 220} 221 222// Print nodes based on the given configuration. 223func (node *Node) Print(opts *Options) { node.print("", opts) } 224 225func dirRecursiveSize(opts *Options, node *Node) (size int64, err error) { 226 if opts.DeepLevel > 0 && node.depth >= opts.DeepLevel { 227 err = errors.New("Depth too high") 228 } 229 230 for _, nnode := range node.nodes { 231 if nnode.err != nil { 232 err = nnode.err 233 continue 234 } 235 236 if !nnode.IsDir() { 237 size += nnode.Size() 238 } else { 239 nsize, e := dirRecursiveSize(opts, nnode) 240 size += nsize 241 if e != nil { 242 err = e 243 } 244 } 245 } 246 return 247} 248 249func (node *Node) print(indent string, opts *Options) { 250 if node.err != nil { 251 err := node.err.Error() 252 if msgs := strings.Split(err, ": "); len(msgs) > 1 { 253 err = msgs[1] 254 } 255 fmt.Printf("%s [%s]\n", node.path, err) 256 return 257 } 258 if !node.IsDir() { 259 var props []string 260 ok, inode, device, uid, gid := getStat(node) 261 // inodes 262 if ok && opts.Inodes { 263 props = append(props, fmt.Sprintf("%d", inode)) 264 } 265 // device 266 if ok && opts.Device { 267 props = append(props, fmt.Sprintf("%3d", device)) 268 } 269 // Mode 270 if opts.FileMode { 271 props = append(props, node.Mode().String()) 272 } 273 // Owner/Uid 274 if ok && opts.ShowUid { 275 uidStr := strconv.Itoa(int(uid)) 276 if u, err := user.LookupId(uidStr); err != nil { 277 props = append(props, fmt.Sprintf("%-8s", uidStr)) 278 } else { 279 props = append(props, fmt.Sprintf("%-8s", u.Username)) 280 } 281 } 282 // Gorup/Gid 283 // TODO: support groupname 284 if ok && opts.ShowGid { 285 gidStr := strconv.Itoa(int(gid)) 286 props = append(props, fmt.Sprintf("%-4s", gidStr)) 287 } 288 // Size 289 if opts.ByteSize || opts.UnitSize { 290 var size string 291 if opts.UnitSize { 292 size = fmt.Sprintf("%4s", formatBytes(node.Size())) 293 } else { 294 size = fmt.Sprintf("%11d", node.Size()) 295 } 296 props = append(props, size) 297 } 298 // Last modification 299 if opts.LastMod { 300 props = append(props, node.ModTime().Format("Jan 02 15:04")) 301 } 302 // Print properties 303 if len(props) > 0 { 304 fmt.Fprintf(opts.OutFile, "[%s] ", strings.Join(props, " ")) 305 } 306 } else { 307 var props []string 308 // Size 309 if opts.ByteSize || opts.UnitSize { 310 var size string 311 rsize, err := dirRecursiveSize(opts, node) 312 if err != nil && rsize <= 0 { 313 if opts.UnitSize { 314 size = "????" 315 } else { 316 size = "???????????" 317 } 318 } else if opts.UnitSize { 319 size = fmt.Sprintf("%4s", formatBytes(rsize)) 320 } else { 321 size = fmt.Sprintf("%11d", rsize) 322 } 323 props = append(props, size) 324 } 325 // Print properties 326 if len(props) > 0 { 327 fmt.Fprintf(opts.OutFile, "[%s] ", strings.Join(props, " ")) 328 } 329 } 330 // name/path 331 var name string 332 if node.depth == 0 || opts.FullPath { 333 name = node.path 334 } else { 335 name = node.Name() 336 } 337 // Quotes 338 if opts.Quotes { 339 name = fmt.Sprintf("\"%s\"", name) 340 } 341 // Colorize 342 if opts.Colorize { 343 name = opts.color(node, name) 344 } 345 // IsSymlink 346 if node.Mode()&os.ModeSymlink == os.ModeSymlink { 347 vtarget, err := os.Readlink(node.path) 348 if err != nil { 349 vtarget = node.path 350 } 351 targetPath, err := filepath.EvalSymlinks(node.path) 352 if err != nil { 353 targetPath = vtarget 354 } 355 fi, err := opts.Fs.Stat(targetPath) 356 if opts.Colorize && fi != nil { 357 vtarget = opts.color(&Node{FileInfo: fi, path: vtarget}, vtarget) 358 } 359 name = fmt.Sprintf("%s -> %s", name, vtarget) 360 // Follow symbolic links like directories 361 if opts.FollowLink { 362 path, err := filepath.Abs(targetPath) 363 if err == nil && fi != nil && fi.IsDir() { 364 if _, ok := node.vpaths[filepath.Clean(path)]; !ok { 365 inf := &Node{FileInfo: fi, path: targetPath} 366 inf.vpaths = node.vpaths 367 inf.Visit(opts) 368 node.nodes = inf.nodes 369 } else { 370 name += " [recursive, not followed]" 371 } 372 } 373 } 374 } 375 // Print file details 376 // the main idea of the print logic came from here: github.com/campoy/tools/tree 377 fmt.Fprintln(opts.OutFile, name) 378 add := "│ " 379 for i, nnode := range node.nodes { 380 if opts.NoIndent { 381 add = "" 382 } else { 383 if i == len(node.nodes)-1 { 384 fmt.Fprintf(opts.OutFile, indent+"└── ") 385 add = " " 386 } else { 387 fmt.Fprintf(opts.OutFile, indent+"├── ") 388 } 389 } 390 nnode.print(indent+add, opts) 391 } 392} 393 394const ( 395 _ = iota // ignore first value by assigning to blank identifier 396 KB int64 = 1 << (10 * iota) 397 MB 398 GB 399 TB 400 PB 401 EB 402) 403 404// Convert bytes to human readable string. Like a 2 MB, 64.2 KB, 52 B 405func formatBytes(i int64) (result string) { 406 var n float64 407 sFmt, eFmt := "%.01f", "" 408 switch { 409 case i > EB: 410 eFmt = "E" 411 n = float64(i) / float64(EB) 412 case i > PB: 413 eFmt = "P" 414 n = float64(i) / float64(PB) 415 case i > TB: 416 eFmt = "T" 417 n = float64(i) / float64(TB) 418 case i > GB: 419 eFmt = "G" 420 n = float64(i) / float64(GB) 421 case i > MB: 422 eFmt = "M" 423 n = float64(i) / float64(MB) 424 case i > KB: 425 eFmt = "K" 426 n = float64(i) / float64(KB) 427 default: 428 sFmt = "%.0f" 429 n = float64(i) 430 } 431 if eFmt != "" && n >= 10 { 432 sFmt = "%.0f" 433 } 434 result = fmt.Sprintf(sFmt+eFmt, n) 435 result = strings.Trim(result, " ") 436 return 437} 438