package main import ( "errors" "fmt" "path/filepath" "strings" ) const ( TopType = "top" IntermediateType = "intermediate" BottomType = "bottom" SingleType = "single" ) var ErrUnknownType = errors.New("unknown type") var ErrShouldBeBottom = errors.New("should be of type Bottom") var ErrChannelError = errors.New("channel errored") var ( NodeB = "│" NodeT = "├──" NodeL = "└──" NodePad = " " ) type Node interface { Size() int64 Spawn(int) error Collect() error Collector() chan int64 Level() int } func isLastLevel(node Node, maxLevel int) bool { return maxLevel-1 == node.Level() } func isBeforeLastLevel(node Node, maxLevel int) bool { return maxLevel-1 > node.Level() } ////////////////////////////// /////////// Top ////////////// ////////////////////////////// type Top struct { path string size int64 tree []Node collect chan int64 level int } func NewTop(path string) *Top { t := &Top{path: path, level: 0} t.collect = make(chan int64) return t } func (t *Top) Size() int64 { return t.size } func (t *Top) Spawn(maxLevel int) error { if t.Level() >= maxLevel { return ErrShouldBeBottom } files, err := ListDir(t.path) if err != nil { return err } for _, info := range files { switch mode := info.Mode(); { case mode.IsDir(): var child Node path := filepath.Join(t.path, info.Name()) if isLastLevel(t, maxLevel) || !AnyDirectoryDownThere(path) { child = t.newBottom(path) } else { child = t.newIntermediate(path) } go child.Spawn(maxLevel) defer func() { go child.Collect() }() case !mode.IsDir() && mode.IsRegular(): t.collect <- info.Size() } } return nil } func (t *Top) Collect() error { for s := range t.collect { t.size += s } return nil } func (t *Top) Collector() chan int64 { return t.collect } func (t *Top) Level() int { return t.level } func (t *Top) Name() string { return filepath.Base(t.path) } func (t *Top) String() string { var out string var lines []string out += fmt.Sprintf("(%d) %s\n", t.Size(), t.path) treeSize := len(t.tree) for _, child := range t.tree { lines = append(lines, fmt.Sprintf("%s", child)) } for a := 0; a < treeSize-1; a++ { childLines := strings.Split(lines[a], "\n") out += fmt.Sprintf(" %s%s\n", NodeT, childLines[0]) for b := 1; b < len(childLines); b++ { out += fmt.Sprintf(" %s%s\n", NodeB, childLines[b]) } } childLines := strings.Split(lines[treeSize-1], "\n") out += fmt.Sprintf(" %s%s\n", NodeL, childLines[0]) for a := 1; a < len(childLines); a++ { out += fmt.Sprintf(" %s%s\n", NodePad, childLines[a]) } return out } func (t *Top) addToTree(child Node) { t.tree = append(t.tree, child) } func (t *Top) newChild(path, kind string) Node { switch { case kind == IntermediateType: i := NewIntermediate(path, t) t.addToTree(i) return i case kind == BottomType: b := NewBottom(path, t) t.addToTree(b) return b default: panic(ErrUnknownType) } } func (t *Top) newIntermediate(path string) *Intermediate { return t.newChild(path, IntermediateType).(*Intermediate) } func (t *Top) newBottom(path string) *Bottom { return t.newChild(path, BottomType).(*Bottom) } ////////////////////////////// /////// Intermediate ///////// ////////////////////////////// type Intermediate struct { path string size int64 parent Node tree []Node collect chan int64 refer chan int64 level int } func NewIntermediate(path string, parent Node) *Intermediate { i := &Intermediate{path: path, refer: parent.Collector(), parent: parent} i.collect = make(chan int64) i.level = parent.Level() + 1 return i } func (i *Intermediate) Size() int64 { return i.size } func (i *Intermediate) Spawn(maxLevel int) error { if i.Level() >= maxLevel { return ErrShouldBeBottom } files, err := ListDir(i.path) if err != nil { return err } for _, info := range files { switch mode := info.Mode(); { case mode.IsDir(): var child Node path := filepath.Join(i.path, info.Name()) if isLastLevel(i, maxLevel) || !AnyDirectoryDownThere(path) { child = i.newBottom(path) } else { child = i.newIntermediate(path) } go child.Spawn(maxLevel) defer func() { go child.Collect() }() case !mode.IsDir() && mode.IsRegular(): i.collect <- info.Size() } } return nil } func (i *Intermediate) Collect() error { for s := range i.collect { i.size += s i.refer <- i.size } close(i.refer) return nil } func (i *Intermediate) Collector() chan int64 { return i.collect } func (i *Intermediate) Level() int { return i.level } func (i *Intermediate) Name() string { return filepath.Base(i.path) } func (i *Intermediate) String() string { var lines []string out := fmt.Sprintf("(%d) %s\n", i.Size(), i.Name()) treeSize := len(i.tree) for _, child := range i.tree { lines = append(lines, fmt.Sprintf("%s", child)) } for a := 0; a < treeSize-1; a++ { childLines := strings.Split(lines[a], "\n") out += fmt.Sprintf("%s%s%s\n", NodePad, NodeT, childLines[0]) for b := 1; b < len(childLines); b++ { out += fmt.Sprintf("%s%s%s\n", NodePad, NodeB, childLines[b]) } } childLines := strings.Split(lines[treeSize-1], "\n") lenChildLines := len(childLines) out += fmt.Sprintf("%s%s%s\n", NodePad, NodeL, childLines[0]) for a := 1; a < lenChildLines-2; a++ { out += fmt.Sprintf("%s%s%s\n", NodePad, NodePad, childLines[a]) } if lenChildLines > 0 { out += fmt.Sprintf("%s%s%s", NodePad, NodePad, childLines[lenChildLines-1]) } return out } func (i *Intermediate) addToTree(child Node) { i.tree = append(i.tree, child) } func (i *Intermediate) newChild(path, kind string) Node { switch { case kind == IntermediateType: c := NewIntermediate(path, i) i.addToTree(c) return c case kind == BottomType: b := NewBottom(path, i) i.addToTree(b) return b default: panic(ErrUnknownType) } } func (i *Intermediate) newIntermediate(path string) *Intermediate { return i.newChild(path, IntermediateType).(*Intermediate) } func (i *Intermediate) newBottom(path string) *Bottom { return i.newChild(path, BottomType).(*Bottom) } ////////////////////////////// ////////// Bottom //////////// ////////////////////////////// type Bottom struct { path string size int64 parent Node walker Walker collect chan int64 refer chan int64 level int } func NewBottom(path string, parent Node) *Bottom { b := &Bottom{path: path, refer: parent.Collector(), parent: parent} b.collect = make(chan int64) b.level = parent.Level() + 1 return b } func (b *Bottom) Size() int64 { return b.size } func (b *Bottom) Spawn(maxLevel int) error { b.walker = NewRealWalker(b.path, b.collect) go b.walker.Walk() return nil } func (b *Bottom) Collect() error { for s := range b.collect { b.size += s b.refer <- b.size } close(b.refer) return nil } func (b *Bottom) Collector() chan int64 { return b.collect } func (b *Bottom) Level() int { return b.level } func (b *Bottom) Name() string { return filepath.Base(b.path) } func (b *Bottom) String() string { return fmt.Sprintf("(%d) %s", b.Size(), b.Name()) } ////////////////////////////// ////////// Single //////////// ////////////////////////////// type Single struct { path string size int64 walker Walker collect chan int64 } func (s *Single) Size() int64 { return s.size } func (s *Single) Spawn(maxLevel int) error { s.walker = NewRealWalker(s.path, s.collect) go s.walker.Walk() return nil } func (s *Single) Collect() error { for v := range s.collect { s.size += v } return nil } func (s *Single) Collector() chan int64 { return s.collect } func (s *Single) Level() int { return 0 } func NewSingle(path string) *Single { s := &Single{path: path} s.collect = make(chan int64) return s } ////////////////////////////// ////////// Walker //////////// ////////////////////////////// type Walker interface { Walk() }