Make the command stop when finished

This commit is contained in:
Blallo 2020-03-07 17:23:06 +01:00
parent 6ccb16e7bd
commit d00a7fa2c0
No known key found for this signature in database
GPG key ID: 0CBE577C9B72DC3F
4 changed files with 251 additions and 153 deletions

View file

@ -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
}
}
}
}

View file

@ -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
}
//////////////////////////////
@ -105,13 +46,13 @@ type Top struct {
size int64
unit ByteSize
tree []Node
collect chan int64
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
}
//////////////////////////////
@ -236,17 +190,19 @@ type Intermediate struct {
unit ByteSize
parent Node
tree []Node
collect chan int64
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
}
//////////////////////////////
@ -384,14 +352,15 @@ type Bottom struct {
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())
}
@ -444,8 +420,25 @@ func (b *Bottom) String() string {
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 {
@ -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())
}
//////////////////////////////

View file

@ -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
View 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
}