Refactor of Ruspa and new Trivella to print until fixed depth.

main.go

@@ -9,11 +9,13 @@ import (
 func main() {
 	var path string
 	var debug bool
+	var depth uint
 	flag.StringVar(&path, "path", ".", "The base path.")
 	flag.BoolVar(&debug, "debug", false, "Toggle debug messages.")
+	flag.UintVar(&depth, "depth", 1, "Max depth of the nodes to display.")
 	flag.Parse()
 	sizer.Debug = debug
 	// rootINode := sizer.NewINode(sizer.FileTypes["dir"], 0, path)
 	// rootINode.Sizer()
-	sizer.NastroConvogliatore(path)
+	sizer.NastroConvogliatore(path, depth)
 }

sizer/worker.go

@@ -1,124 +1,157 @@
 package sizer
 
 import (
+	"bufio"
 	"fmt"
-
-	"github.com/gosuri/uilive"
+	"io"
+	"os"
 )
 
+type dummy struct{}
+
+func (d dummy) Write(p []byte) (n int, err error) {
+	return 0, nil
+}
+
+// Trivella provides an io.Writer with a notion of depth in the
+// path graph
+type Trivella struct {
+	Level  uint
+	Writer io.Writer
+}
+
+// Scendi increases the depth in the graph
+func (t *Trivella) Scendi() *Trivella {
+	var v *Trivella
+	if t.Level > 0 {
+		v = TrInit(t.Writer, t.Level-1)
+	} else {
+		var d dummy
+		v = TrInit(d, 0)
+	}
+	return v
+}
+
+// TrInit assigns the writer to a Trivella
+func TrInit(writer io.Writer, level uint) *Trivella {
+	var t Trivella
+	t.Writer = writer
+	t.Level = level
+	return &t
+}
+
 // Ruspa is the asynchronous worker that is in charge of collecting
 // the underlying sizes and relay the total to the eventual parent.
 type Ruspa struct {
-	INode  *INode
-	Send   chan StatusReport
-	Recv   chan StatusReport
-	Done   chan StatusReport
-	Opened map[string]*Ruspa
+	INode     *INode
+	ReportIn  chan StatusReport
+	ReportOut chan StatusReport
+	DoneIn    chan StatusReport
+	DoneOut   chan StatusReport
+	Opened    map[string]*Ruspa
 }
 
 // NewRuspa inits a new Ruspa
-func NewRuspa(inode *INode, collect chan StatusReport) *Ruspa {
+func NewRuspa(inode *INode, reportOut, doneOut chan StatusReport) *Ruspa {
 	var r Ruspa
 	r.INode = inode
-	r.Recv = make(chan StatusReport)
-	if collect == nil {
-		r.Send = make(chan StatusReport)
-	} else {
-		r.Send = collect
-	}
-	r.Done = make(chan StatusReport)
+	r.ReportOut = reportOut
+	r.DoneOut = doneOut
+	r.ReportIn = make(chan StatusReport)
+	r.DoneIn = make(chan StatusReport)
 	r.Opened = make(map[string]*Ruspa)
 	return &r
 }
 
+// Ammucchia gathers the size of the underlying INode(s) from the Ruspa worker
+// and adds to the INode size.
+func (w *Ruspa) Ammucchia(t *Trivella) {
+	Console.Debugln(Gray("Ammucchia:", w.INode.Path))
+	for len(w.Opened) > 0 {
+		select {
+		case report := <-w.ReportIn:
+			Console.Debugln(Gray("ReportIn:", report))
+			if report.err == nil {
+				w.INode.Size += report.size
+				fmt.Fprintf(t.Writer, "%s: %d\n", w.INode.Path, w.INode.Size)
+			} else {
+				Console.Debugln(Red(report.err))
+			}
+		case leaf := <-w.DoneIn:
+			Console.Debugln(Gray("Done:", leaf.path))
+			delete(w.Opened, leaf.path)
+		}
+	}
+	Console.Debugln(Gray(w.INode.Path, ": Completed. Exiting..."))
+	w.DoneOut <- StatusReport{path: w.INode.Path, size: w.INode.Size}
+}
+
 // Scava starts a worker on the given INode.
-func (w *Ruspa) Scava(done chan StatusReport) {
+func (w *Ruspa) Scava(t *Trivella) {
 	Console.Debugln(Gray("Scava: ", w.INode.Path))
 	children, err := ls(w.INode.Path)
 	if err != nil {
-		w.Send <- StatusReport{path: w.INode.Path, err: err}
+		w.ReportOut <- StatusReport{path: w.INode.Path, err: err}
 	}
 	for _, childPath := range children {
 		kind, size, err := IdentifyType(childPath)
 		if err != nil {
-			w.Send <- StatusReport{path: childPath, err: err}
+			w.ReportOut <- StatusReport{path: childPath, err: err}
 		}
 		childINode := NewINode(kind, size, childPath)
 		w.INode.Children.Append(childINode)
 		switch {
 		case kind == FileType:
 			Console.Debugln(Gray("[file]", w.INode.Path, " ~> ", childINode.Path))
-			w.Send <- StatusReport{path: childPath, size: size}
+			w.ReportOut <- StatusReport{path: childPath, size: size}
 		case kind == DirType:
 			Console.Debugln(Gray("[dir]", w.INode.Path, " ~> ", childINode.Path))
-			cw := NewRuspa(childINode, w.Recv)
+			w.INode.Size += size
+			cw := NewRuspa(childINode, w.ReportIn, w.DoneIn)
 			w.Opened[childPath] = cw
-			go cw.Scava(w.Done)
-			// go cw.Ammucchia(w.Done)
+			go cw.Scava(t.Scendi())
 		}
 	}
-	if len(w.Opened) == 0 {
-		done <- StatusReport{path: w.INode.Path}
+	if len(w.Opened) != 0 {
+		Console.Debugln(Gray("Spawning Ammucchia:", w.INode.Path))
+		go w.Ammucchia(t)
 	} else {
-		go w.Ammucchia(done)
-	}
-}
-
-// Ammucchia gathers the size of the underlying Ruspa workers
-// and adds to the INode size.
-func (w *Ruspa) Ammucchia(done chan StatusReport) {
-	for len(w.Opened) > 0 {
-		select {
-		case report := <-w.Recv:
-			if report.err != nil {
-				Console.Debugln(Gray(Red(report.err)))
-			}
-			if report.err == nil {
-				w.INode.Size += report.size
-			}
-		case leaf := <-w.Done:
-			Console.Debugln(Gray("Done:", leaf.path))
-			delete(w.Opened, leaf.path)
-		}
+		w.DoneOut <- StatusReport{path: w.INode.Path, size: w.INode.Size}
 	}
-	Console.Debugln(Gray(w.INode.Path, ": Completed. Exiting..."))
-	done <- StatusReport{path: w.INode.Path, size: w.INode.Size}
 }
 
 // Pesa starts the worker for current inode and returns the channel
 // to listen to get the size.
-func (i *INode) Pesa(done chan StatusReport) chan StatusReport {
-	result := make(chan StatusReport)
-	w := NewRuspa(i, result)
-	go w.Scava(done)
-	// go w.Ammucchia(done)
-	return result
+func (i *INode) Pesa(t *Trivella, report, done chan StatusReport) {
+	w := NewRuspa(i, report, done)
+	go w.Scava(t)
 }
 
 // NastroConvogliatore starts a Pesa on a given path and displays
 // the current value of the size.
-func NastroConvogliatore(path string) {
-	var size int64
+func NastroConvogliatore(path string, depth uint) {
+	finished := false
 	pathINode, err := INodeFromPath(path)
 	if err != nil {
 		Console.Fatalln("Failed creating inode:", err)
 	}
 	Console.Debugln(Gray("Starting Pesa on path:", path))
+	report := make(chan StatusReport)
 	done := make(chan StatusReport)
-	result := pathINode.Pesa(done)
-	writer := uilive.New()
-	writer.Start()
+	writer := bufio.NewWriter(os.Stdout)
+	defer writer.Flush()
+	t := TrInit(writer, depth)
 	fmt.Fprintf(writer, "%s\n", Green("Starting..."))
-	size = 0
-	for {
+	pathINode.Pesa(t, report, done)
+	for !finished {
 		select {
-		case res := <-result:
-			size += res.size
-			fmt.Fprintf(writer, "%s: %d\n%s: %d\n", res.path, res.size, path, size)
-		case <-done:
-			fmt.Fprintf(writer, "%s: %d\n", path, pathINode.Size)
-			writer.Stop()
-			return
+		case res := <-report:
+			Console.Debugln(Gray("MAIN: result =>", res))
+		case res := <-done:
+			Console.Debugln(Gray("MAIN: done =>", res))
+			finished = true
 		}
 	}
+	fmt.Fprintf(writer, "%s: %d\n", path, pathINode.Size)
+	return
 }