514 rivejä
10 KiB
Go
514 rivejä
10 KiB
Go
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()
|
|
}
|