package tree 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 ErrUnknownUnit = errors.New("unknown unit") var ( NodeB = "│" NodeT = "├──" NodeL = "└──" NodePad = " " ) type Node interface { SetUnit(string) GetUnit() ByteSize Size() int64 Spawn(int) error Collect() error AddCollector(chan int64) Level() int Name() string Complete() bool Depth() int } ////////////////////////////// /////////// Top ////////////// ////////////////////////////// type Top struct { path string size int64 unit ByteSize tree []Node collect []<-chan int64 level int complete bool } func NewTop(path string) *Top { t := &Top{path: path, level: 0, unit: KB} return t } func (t *Top) SetUnit(unit string) { t.unit = setUnit(unit) } func (t *Top) GetUnit() ByteSize { return t.unit } 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 var collect chan int64 path := filepath.Join(t.path, info.Name()) if isLastLevel(t, maxLevel) || !AnyDirectoryDownThere(path) { child, collect = t.newBottom(path) } else { child, collect = t.newIntermediate(path) } t.AddCollector(collect) go child.Spawn(maxLevel) case !mode.IsDir() && mode.IsRegular(): t.size += info.Size() } } go t.Collect() return nil } func (t *Top) Collect() error { for s := range merge(t.collect) { t.size += s } t.complete = true return nil } func (t *Top) AddCollector(collect chan int64) { t.collect = append(t.collect, collect) } func (t *Top) Level() int { return t.level } func (t *Top) Name() string { return filepath.Base(t.path) } func (t *Top) Complete() bool { for _, child := range t.tree { if !child.Complete() { return false } } return true } func (t *Top) Depth() int { var depth int for _, child := range t.tree { depth += child.Depth() } return depth } func (t *Top) String() string { var out string var lines []string out += fmt.Sprintf("(%s) %s\n", fmtSize(t.Size(), t.GetUnit()), 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", " ", childLines[a]) } return out } func (t *Top) addToTree(child Node) { t.tree = append(t.tree, child) } func (t *Top) newChild(path, kind string) (Node, chan int64) { switch { case kind == IntermediateType: i, collect := NewIntermediate(path, t) t.addToTree(i) return i, collect case kind == BottomType: b, collect := NewBottom(path, t) t.addToTree(b) return b, collect default: panic(ErrUnknownType) } } func (t *Top) newIntermediate(path string) (*Intermediate, chan int64) { child, collect := t.newChild(path, IntermediateType) return child.(*Intermediate), collect } func (t *Top) newBottom(path string) (*Bottom, chan int64) { child, collect := t.newChild(path, BottomType) return child.(*Bottom), collect } ////////////////////////////// /////// Intermediate ///////// ////////////////////////////// type Intermediate struct { path string size int64 unit ByteSize parent Node tree []Node collect []<-chan int64 refer chan int64 level int complete bool } func NewIntermediate(path string, parent Node) (*Intermediate, chan int64) { refer := make(chan int64) i := &Intermediate{path: path, refer: refer, parent: parent} i.collect = make([]<-chan int64, 1) i.level = parent.Level() + 1 i.unit = parent.GetUnit() return i, refer } func (i *Intermediate) SetUnit(unit string) { i.unit = setUnit(unit) } func (i *Intermediate) GetUnit() ByteSize { return i.unit } 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 var collect chan int64 path := filepath.Join(i.path, info.Name()) if isLastLevel(i, maxLevel) || !AnyDirectoryDownThere(path) { child, collect = i.newBottom(path) } else { child, collect = i.newIntermediate(path) } i.AddCollector(collect) go child.Spawn(maxLevel) case !mode.IsDir() && mode.IsRegular(): i.size += info.Size() i.refer <- info.Size() } } go i.Collect() return nil } func (i *Intermediate) Collect() error { for s := range merge(i.collect) { i.size += s i.refer <- s } i.complete = true return nil } func (i *Intermediate) AddCollector(collect chan int64) { i.collect = append(i.collect, collect) } func (i *Intermediate) Level() int { return i.level } func (i *Intermediate) Name() string { return filepath.Base(i.path) } func (i *Intermediate) Complete() bool { for _, child := range i.tree { if !child.Complete() { return false } } return true } func (i *Intermediate) Depth() int { var depth int for _, child := range i.tree { depth += child.Depth() } return depth } func (i *Intermediate) String() string { var lines []string out := fmt.Sprintf("(%s) %s\n", fmtSize(i.Size(), i.GetUnit()), 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) if lenChildLines > 1 { out += fmt.Sprintf("%s%s%s\n", NodePad, NodeL, childLines[0]) for a := 1; a < lenChildLines-2; a++ { out += fmt.Sprintf(" %s%s\n", NodePad, childLines[a]) } out += fmt.Sprintf(" %s%s", NodePad, childLines[lenChildLines-1]) } else { out += fmt.Sprintf("%s%s%s", NodePad, NodeL, childLines[0]) } return out } func (i *Intermediate) addToTree(child Node) { i.tree = append(i.tree, child) } func (i *Intermediate) newChild(path, kind string) (Node, chan int64) { switch { case kind == IntermediateType: c, collect := NewIntermediate(path, i) i.addToTree(c) return c, collect case kind == BottomType: b, collect := NewBottom(path, i) i.addToTree(b) return b, collect default: panic(ErrUnknownType) } } func (i *Intermediate) newIntermediate(path string) (*Intermediate, chan int64) { child, collect := i.newChild(path, IntermediateType) return child.(*Intermediate), collect } func (i *Intermediate) newBottom(path string) (*Bottom, chan int64) { child, collect := i.newChild(path, BottomType) return child.(*Bottom), collect } ////////////////////////////// ////////// Bottom //////////// ////////////////////////////// type Bottom struct { path string size int64 unit ByteSize parent Node walker Walker collect chan int64 refer chan int64 level int complete bool } func NewBottom(path string, parent Node) (*Bottom, chan int64) { refer := make(chan int64) b := &Bottom{path: path, refer: refer, parent: parent, unit: parent.GetUnit()} b.level = parent.Level() + 1 b.unit = parent.GetUnit() return b, refer } func (b *Bottom) SetUnit(unit string) { b.unit = setUnit(unit) } func (b *Bottom) GetUnit() ByteSize { return b.unit } func (b *Bottom) Size() int64 { return b.size } func (b *Bottom) Spawn(maxLevel int) error { collect := make(chan int64) b.AddCollector(collect) b.walker = NewRealWalker(b.path, collect) go b.walker.Walk() go b.Collect() return nil } func (b *Bottom) Collect() error { for s := range b.collect { b.size += s b.refer <- s } b.complete = true return nil } func (b *Bottom) AddCollector(collect chan int64) { b.collect = collect } func (b *Bottom) Level() int { return b.level } func (b *Bottom) Name() string { return filepath.Base(b.path) } func (b *Bottom) Complete() bool { return b.complete } func (b *Bottom) Depth() int { return 1 } func (b *Bottom) String() string { return fmt.Sprintf("(%s) %s", fmtSize(b.Size(), b.GetUnit()), b.Name()) } ////////////////////////////// ////////// Single //////////// ////////////////////////////// type Single struct { path string size int64 unit ByteSize walker Walker collect chan int64 complete bool } func NewSingle(path string) *Single { s := &Single{path: path} collect := make(chan int64) s.AddCollector(collect) return s } func (s *Single) SetUnit(unit string) { s.unit = setUnit(unit) } func (s *Single) GetUnit() ByteSize { return s.unit } 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() go s.Collect() return nil } func (s *Single) Collect() error { for v := range s.collect { s.size += v } s.complete = true return nil } func (s *Single) AddCollector(collect chan int64) { s.collect = collect } func (s *Single) Level() int { return 0 } func (s *Single) Name() string { return filepath.Base(s.path) } func (s *Single) Complete() bool { return s.complete } func (s *Single) Depth() int { return 1 } func (s *Single) String() string { return fmt.Sprintf("(%s) %s\n", fmtSize(s.Size(), s.GetUnit()), s.Name()) } ////////////////////////////// ////////// Walker //////////// ////////////////////////////// type Walker interface { Walk() }