Make the command stop when finished
This commit is contained in:
parent
6ccb16e7bd
commit
d00a7fa2c0
4 changed files with 251 additions and 153 deletions
|
|
@ -44,26 +44,29 @@ func (u *UnitValue) Set(value string) error {
|
|||
func main() {
|
||||
var path string
|
||||
var depth int
|
||||
var t tree.Node
|
||||
var root tree.Node
|
||||
var unit = &UnitValue{unit: "KB"}
|
||||
flag.StringVar(&path, "path", ".", "Path from where to start the walk from")
|
||||
flag.IntVar(&depth, "depth", 0, "Depth to display")
|
||||
flag.Var(unit, "unit", "Unit in which to report size")
|
||||
flag.Parse()
|
||||
|
||||
if depth == 0 {
|
||||
t = tree.NewSingle(path)
|
||||
if depth == 0 || !tree.AnyDirectoryDownThere(path) {
|
||||
root = tree.NewSingle(path)
|
||||
} else {
|
||||
t = tree.NewTop(path)
|
||||
root = tree.NewTop(path)
|
||||
}
|
||||
t.SetUnit(unit.String())
|
||||
go t.Spawn(depth)
|
||||
go t.Collect()
|
||||
for !t.Complete() {
|
||||
root.SetUnit(unit.String())
|
||||
go root.Spawn(depth)
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
fmt.Printf("\033[H\033[2J")
|
||||
fmt.Println(t)
|
||||
fmt.Println(root)
|
||||
if root.Complete() {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,17 +14,6 @@ const (
|
|||
SingleType = "single"
|
||||
)
|
||||
|
||||
type ByteSize uint
|
||||
|
||||
const (
|
||||
B ByteSize = 1 << (10 * iota)
|
||||
KB
|
||||
MB
|
||||
GB
|
||||
TB
|
||||
PB
|
||||
)
|
||||
|
||||
var ErrUnknownType = errors.New("unknown type")
|
||||
var ErrShouldBeBottom = errors.New("should be of type Bottom")
|
||||
var ErrChannelError = errors.New("channel errored")
|
||||
|
|
@ -42,58 +31,10 @@ type Node interface {
|
|||
Size() int64
|
||||
Spawn(int) error
|
||||
Collect() error
|
||||
Collector() chan int64
|
||||
AddCollector(chan int64)
|
||||
Level() int
|
||||
}
|
||||
|
||||
func isLastLevel(node Node, maxLevel int) bool {
|
||||
return maxLevel-1 == node.Level()
|
||||
}
|
||||
|
||||
func setUnit(unit string) ByteSize {
|
||||
switch unit {
|
||||
case "B":
|
||||
return B
|
||||
case "KB":
|
||||
return KB
|
||||
case "MB":
|
||||
return MB
|
||||
case "GB":
|
||||
return GB
|
||||
case "TB":
|
||||
return TB
|
||||
case "PB":
|
||||
// I reckon this will never happen, but you never know
|
||||
return PB
|
||||
default:
|
||||
panic(ErrUnknownUnit)
|
||||
}
|
||||
}
|
||||
|
||||
func fmtUnit(unit ByteSize) string {
|
||||
switch unit {
|
||||
case B:
|
||||
return "B"
|
||||
case KB:
|
||||
return "KB"
|
||||
case MB:
|
||||
return "MB"
|
||||
case GB:
|
||||
return "GB"
|
||||
case TB:
|
||||
return "TB"
|
||||
case PB:
|
||||
// I reckon this will never happen, but you never know
|
||||
return "PB"
|
||||
default:
|
||||
panic(ErrUnknownUnit)
|
||||
}
|
||||
}
|
||||
|
||||
func fmtSize(size int64, unit ByteSize) string {
|
||||
var dimension float64
|
||||
dimension = float64(size) / float64(unit)
|
||||
return fmt.Sprintf("%.2f %s", dimension, fmtUnit(unit))
|
||||
Name() string
|
||||
Complete() bool
|
||||
}
|
||||
|
||||
//////////////////////////////
|
||||
|
|
@ -101,17 +42,17 @@ func fmtSize(size int64, unit ByteSize) string {
|
|||
//////////////////////////////
|
||||
|
||||
type Top struct {
|
||||
path string
|
||||
size int64
|
||||
unit ByteSize
|
||||
tree []Node
|
||||
collect chan int64
|
||||
level int
|
||||
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}
|
||||
t.collect = make(chan int64)
|
||||
return t
|
||||
}
|
||||
|
||||
|
|
@ -140,32 +81,34 @@ func (t *Top) Spawn(maxLevel int) error {
|
|||
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 = t.newBottom(path)
|
||||
child, collect = t.newBottom(path)
|
||||
} else {
|
||||
child = t.newIntermediate(path)
|
||||
child, collect = t.newIntermediate(path)
|
||||
}
|
||||
t.AddCollector(collect)
|
||||
go child.Spawn(maxLevel)
|
||||
defer func() {
|
||||
go child.Collect()
|
||||
}()
|
||||
case !mode.IsDir() && mode.IsRegular():
|
||||
t.collect <- info.Size()
|
||||
t.size += info.Size()
|
||||
}
|
||||
}
|
||||
|
||||
go t.Collect()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Top) Collect() error {
|
||||
for s := range t.collect {
|
||||
for s := range merge(t.collect) {
|
||||
t.size += s
|
||||
}
|
||||
t.complete = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Top) Collector() chan int64 {
|
||||
return t.collect
|
||||
func (t *Top) AddCollector(collect chan int64) {
|
||||
t.collect = append(t.collect, collect)
|
||||
}
|
||||
|
||||
func (t *Top) Level() int {
|
||||
|
|
@ -176,6 +119,15 @@ 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) String() string {
|
||||
var out string
|
||||
var lines []string
|
||||
|
|
@ -203,27 +155,29 @@ func (t *Top) addToTree(child Node) {
|
|||
t.tree = append(t.tree, child)
|
||||
}
|
||||
|
||||
func (t *Top) newChild(path, kind string) Node {
|
||||
func (t *Top) newChild(path, kind string) (Node, chan int64) {
|
||||
switch {
|
||||
case kind == IntermediateType:
|
||||
i := NewIntermediate(path, t)
|
||||
i, collect := NewIntermediate(path, t)
|
||||
t.addToTree(i)
|
||||
return i
|
||||
return i, collect
|
||||
case kind == BottomType:
|
||||
b := NewBottom(path, t)
|
||||
b, collect := NewBottom(path, t)
|
||||
t.addToTree(b)
|
||||
return b
|
||||
return b, collect
|
||||
default:
|
||||
panic(ErrUnknownType)
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Top) newIntermediate(path string) *Intermediate {
|
||||
return t.newChild(path, IntermediateType).(*Intermediate)
|
||||
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 {
|
||||
return t.newChild(path, BottomType).(*Bottom)
|
||||
func (t *Top) newBottom(path string) (*Bottom, chan int64) {
|
||||
child, collect := t.newChild(path, BottomType)
|
||||
return child.(*Bottom), collect
|
||||
}
|
||||
|
||||
//////////////////////////////
|
||||
|
|
@ -231,22 +185,24 @@ func (t *Top) newBottom(path string) *Bottom {
|
|||
//////////////////////////////
|
||||
|
||||
type Intermediate struct {
|
||||
path string
|
||||
size int64
|
||||
unit ByteSize
|
||||
parent Node
|
||||
tree []Node
|
||||
collect chan int64
|
||||
refer chan int64
|
||||
level int
|
||||
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 {
|
||||
i := &Intermediate{path: path, refer: parent.Collector(), parent: parent}
|
||||
i.collect = make(chan int64)
|
||||
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
|
||||
return i, refer
|
||||
}
|
||||
|
||||
func (i *Intermediate) SetUnit(unit string) {
|
||||
|
|
@ -275,35 +231,36 @@ func (i *Intermediate) Spawn(maxLevel int) error {
|
|||
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 = i.newBottom(path)
|
||||
child, collect = i.newBottom(path)
|
||||
} else {
|
||||
child = i.newIntermediate(path)
|
||||
child, collect = i.newIntermediate(path)
|
||||
}
|
||||
i.AddCollector(collect)
|
||||
go child.Spawn(maxLevel)
|
||||
defer func() {
|
||||
go child.Collect()
|
||||
}()
|
||||
case !mode.IsDir() && mode.IsRegular():
|
||||
i.collect <- info.Size()
|
||||
i.size += info.Size()
|
||||
i.refer <- info.Size()
|
||||
}
|
||||
}
|
||||
|
||||
go i.Collect()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (i *Intermediate) Collect() error {
|
||||
for s := range i.collect {
|
||||
for s := range merge(i.collect) {
|
||||
i.size += s
|
||||
i.refer <- s
|
||||
}
|
||||
close(i.refer)
|
||||
i.complete = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func (i *Intermediate) Collector() chan int64 {
|
||||
return i.collect
|
||||
func (i *Intermediate) AddCollector(collect chan int64) {
|
||||
i.collect = append(i.collect, collect)
|
||||
}
|
||||
|
||||
func (i *Intermediate) Level() int {
|
||||
|
|
@ -314,6 +271,15 @@ 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) String() string {
|
||||
var lines []string
|
||||
out := fmt.Sprintf("(%s) %s\n", fmtSize(i.Size(), i.GetUnit()), i.Name())
|
||||
|
|
@ -348,27 +314,29 @@ func (i *Intermediate) addToTree(child Node) {
|
|||
i.tree = append(i.tree, child)
|
||||
}
|
||||
|
||||
func (i *Intermediate) newChild(path, kind string) Node {
|
||||
func (i *Intermediate) newChild(path, kind string) (Node, chan int64) {
|
||||
switch {
|
||||
case kind == IntermediateType:
|
||||
c := NewIntermediate(path, i)
|
||||
c, collect := NewIntermediate(path, i)
|
||||
i.addToTree(c)
|
||||
return c
|
||||
return c, collect
|
||||
case kind == BottomType:
|
||||
b := NewBottom(path, i)
|
||||
b, collect := NewBottom(path, i)
|
||||
i.addToTree(b)
|
||||
return b
|
||||
return b, collect
|
||||
default:
|
||||
panic(ErrUnknownType)
|
||||
}
|
||||
}
|
||||
|
||||
func (i *Intermediate) newIntermediate(path string) *Intermediate {
|
||||
return i.newChild(path, IntermediateType).(*Intermediate)
|
||||
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 {
|
||||
return i.newChild(path, BottomType).(*Bottom)
|
||||
func (i *Intermediate) newBottom(path string) (*Bottom, chan int64) {
|
||||
child, collect := i.newChild(path, BottomType)
|
||||
return child.(*Bottom), collect
|
||||
}
|
||||
|
||||
//////////////////////////////
|
||||
|
|
@ -376,22 +344,23 @@ func (i *Intermediate) newBottom(path string) *Bottom {
|
|||
//////////////////////////////
|
||||
|
||||
type Bottom struct {
|
||||
path string
|
||||
size int64
|
||||
unit ByteSize
|
||||
parent Node
|
||||
walker Walker
|
||||
collect chan int64
|
||||
refer chan int64
|
||||
level int
|
||||
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 {
|
||||
b := &Bottom{path: path, refer: parent.Collector(), parent: parent, unit: parent.GetUnit()}
|
||||
b.collect = make(chan int64)
|
||||
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
|
||||
return b, refer
|
||||
}
|
||||
|
||||
func (b *Bottom) SetUnit(unit string) {
|
||||
|
|
@ -407,8 +376,11 @@ func (b *Bottom) Size() int64 {
|
|||
}
|
||||
|
||||
func (b *Bottom) Spawn(maxLevel int) error {
|
||||
b.walker = NewRealWalker(b.path, b.collect)
|
||||
collect := make(chan int64)
|
||||
b.AddCollector(collect)
|
||||
b.walker = NewRealWalker(b.path, collect)
|
||||
go b.walker.Walk()
|
||||
go b.Collect()
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -417,12 +389,12 @@ func (b *Bottom) Collect() error {
|
|||
b.size += s
|
||||
b.refer <- s
|
||||
}
|
||||
close(b.refer)
|
||||
b.complete = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *Bottom) Collector() chan int64 {
|
||||
return b.collect
|
||||
func (b *Bottom) AddCollector(collect chan int64) {
|
||||
b.collect = collect
|
||||
}
|
||||
|
||||
func (b *Bottom) Level() int {
|
||||
|
|
@ -433,6 +405,10 @@ func (b *Bottom) Name() string {
|
|||
return filepath.Base(b.path)
|
||||
}
|
||||
|
||||
func (b *Bottom) Complete() bool {
|
||||
return b.complete
|
||||
}
|
||||
|
||||
func (b *Bottom) String() string {
|
||||
return fmt.Sprintf("(%s) %s", fmtSize(b.Size(), b.GetUnit()), b.Name())
|
||||
}
|
||||
|
|
@ -442,10 +418,27 @@ func (b *Bottom) String() string {
|
|||
//////////////////////////////
|
||||
|
||||
type Single struct {
|
||||
path string
|
||||
size int64
|
||||
walker Walker
|
||||
collect chan int64
|
||||
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 {
|
||||
|
|
@ -455,6 +448,7 @@ func (s *Single) Size() int64 {
|
|||
func (s *Single) Spawn(maxLevel int) error {
|
||||
s.walker = NewRealWalker(s.path, s.collect)
|
||||
go s.walker.Walk()
|
||||
go s.Collect()
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -462,21 +456,28 @@ func (s *Single) Collect() error {
|
|||
for v := range s.collect {
|
||||
s.size += v
|
||||
}
|
||||
s.complete = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Single) Collector() chan int64 {
|
||||
return s.collect
|
||||
func (s *Single) AddCollector(collect chan int64) {
|
||||
s.collect = collect
|
||||
}
|
||||
|
||||
func (s *Single) Level() int {
|
||||
return 0
|
||||
}
|
||||
|
||||
func NewSingle(path string) *Single {
|
||||
s := &Single{path: path}
|
||||
s.collect = make(chan int64)
|
||||
return s
|
||||
func (s *Single) Name() string {
|
||||
return filepath.Base(s.path)
|
||||
}
|
||||
|
||||
func (s *Single) Complete() bool {
|
||||
return s.complete
|
||||
}
|
||||
|
||||
func (s *Single) String() string {
|
||||
return fmt.Sprintf("(%s) %s\n", fmtSize(s.Size(), s.GetUnit()), s.Name())
|
||||
}
|
||||
|
||||
//////////////////////////////
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ func (r *RealWalker) walkFunc(path string, info os.FileInfo, err error) error {
|
|||
|
||||
func (r *RealWalker) Walk() {
|
||||
filepath.Walk(r.path, r.walkFunc)
|
||||
close(r.report)
|
||||
}
|
||||
|
||||
func NewRealWalker(path string, report chan int64) *RealWalker {
|
||||
|
|
|
|||
93
tree/tree_utils.go
Normal file
93
tree/tree_utils.go
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
package tree
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type ByteSize uint
|
||||
|
||||
const (
|
||||
B ByteSize = 1 << (10 * iota)
|
||||
KB
|
||||
MB
|
||||
GB
|
||||
TB
|
||||
PB
|
||||
)
|
||||
|
||||
func isLastLevel(node Node, maxLevel int) bool {
|
||||
return maxLevel-1 == node.Level()
|
||||
}
|
||||
|
||||
func setUnit(unit string) ByteSize {
|
||||
switch unit {
|
||||
case "B":
|
||||
return B
|
||||
case "KB":
|
||||
return KB
|
||||
case "MB":
|
||||
return MB
|
||||
case "GB":
|
||||
return GB
|
||||
case "TB":
|
||||
return TB
|
||||
case "PB":
|
||||
// I reckon this will never happen, but you never know
|
||||
return PB
|
||||
default:
|
||||
panic(ErrUnknownUnit)
|
||||
}
|
||||
}
|
||||
|
||||
func fmtUnit(unit ByteSize) string {
|
||||
switch unit {
|
||||
case B:
|
||||
return "B"
|
||||
case KB:
|
||||
return "KB"
|
||||
case MB:
|
||||
return "MB"
|
||||
case GB:
|
||||
return "GB"
|
||||
case TB:
|
||||
return "TB"
|
||||
case PB:
|
||||
// I reckon this will never happen, but you never know
|
||||
return "PB"
|
||||
default:
|
||||
panic(ErrUnknownUnit)
|
||||
}
|
||||
}
|
||||
|
||||
func fmtSize(size int64, unit ByteSize) string {
|
||||
var dimension float64
|
||||
dimension = float64(size) / float64(unit)
|
||||
return fmt.Sprintf("%.2f %s", dimension, fmtUnit(unit))
|
||||
}
|
||||
|
||||
func merge(cs []<-chan int64) <-chan int64 {
|
||||
var wg sync.WaitGroup
|
||||
out := make(chan int64)
|
||||
|
||||
// Start an output goroutine for each input channel in cs. output
|
||||
// copies values from c to out until c is closed, then calls wg.Done.
|
||||
output := func(c <-chan int64) {
|
||||
for n := range c {
|
||||
out <- n
|
||||
}
|
||||
wg.Done()
|
||||
}
|
||||
wg.Add(len(cs))
|
||||
for _, c := range cs {
|
||||
go output(c)
|
||||
}
|
||||
|
||||
// Start a goroutine to close out once all the output goroutines are
|
||||
// done. This must start after the wg.Add call.
|
||||
go func() {
|
||||
wg.Wait()
|
||||
close(out)
|
||||
}()
|
||||
return out
|
||||
}
|
||||
Loading…
Reference in a new issue