123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514 |
- 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()
- }
|