Introduce unit for size and add tests

dynamic_tree.go

@@ -14,9 +14,21 @@ 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")
+var ErrUnknownUnit = errors.New("unknown unit")
 var (
 	NodeB   = "│"
 	NodeT   = "├──"
@@ -25,6 +37,8 @@ var (
 )
 
 type Node interface {
+	SetUnit(string)
+	GetUnit() ByteSize
 	Size() int64
 	Spawn(int) error
 	Collect() error
@@ -36,6 +50,50 @@ 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))
 }
 
 //////////////////////////////
@@ -45,17 +103,26 @@ func isLastLevel(node Node, maxLevel int) bool {
 type Top struct {
 	path    string
 	size    int64
+	unit    ByteSize
 	tree    []Node
 	collect chan int64
 	level   int
 }
 
 func NewTop(path string) *Top {
-	t := &Top{path: path, level: 0}
+	t := &Top{path: path, level: 0, unit: KB}
 	t.collect = make(chan int64)
 	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
 }
@@ -112,7 +179,7 @@ func (t *Top) Name() string {
 func (t *Top) String() string {
 	var out string
 	var lines []string
-	out += fmt.Sprintf("(%d) %s\n", t.Size(), t.path)
+	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))
@@ -127,7 +194,7 @@ func (t *Top) String() string {
 	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", NodePad, childLines[a])
+		out += fmt.Sprintf("%s%s\n", "  ", childLines[a])
 	}
 	return out
 }
@@ -166,6 +233,7 @@ func (t *Top) newBottom(path string) *Bottom {
 type Intermediate struct {
 	path    string
 	size    int64
+	unit    ByteSize
 	parent  Node
 	tree    []Node
 	collect chan int64
@@ -177,9 +245,18 @@ func NewIntermediate(path string, parent Node) *Intermediate {
 	i := &Intermediate{path: path, refer: parent.Collector(), parent: parent}
 	i.collect = make(chan int64)
 	i.level = parent.Level() + 1
+	i.unit = parent.GetUnit()
 	return i
 }
 
+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
 }
@@ -239,7 +316,7 @@ func (i *Intermediate) Name() string {
 
 func (i *Intermediate) String() string {
 	var lines []string
-	out := fmt.Sprintf("(%d) %s\n", i.Size(), i.Name())
+	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))
@@ -258,9 +335,9 @@ func (i *Intermediate) String() string {
 	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%s\n", NodePad, NodePad, childLines[a])
+			out += fmt.Sprintf(" %s%s\n", NodePad, childLines[a])
 		}
-		out += fmt.Sprintf("%s%s%s", NodePad, NodePad, childLines[lenChildLines-1])
+		out += fmt.Sprintf(" %s%s", NodePad, childLines[lenChildLines-1])
 	} else {
 		out += fmt.Sprintf("%s%s%s", NodePad, NodeL, childLines[0])
 	}
@@ -301,6 +378,7 @@ func (i *Intermediate) newBottom(path string) *Bottom {
 type Bottom struct {
 	path    string
 	size    int64
+	unit    ByteSize
 	parent  Node
 	walker  Walker
 	collect chan int64
@@ -309,12 +387,21 @@ type Bottom struct {
 }
 
 func NewBottom(path string, parent Node) *Bottom {
-	b := &Bottom{path: path, refer: parent.Collector(), parent: parent}
+	b := &Bottom{path: path, refer: parent.Collector(), parent: parent, unit: parent.GetUnit()}
 	b.collect = make(chan int64)
 	b.level = parent.Level() + 1
+	b.unit = parent.GetUnit()
 	return b
 }
 
+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
 }
@@ -347,7 +434,7 @@ func (b *Bottom) Name() string {
 }
 
 func (b *Bottom) String() string {
-	return fmt.Sprintf("(%d) %s", b.Size(), b.Name())
+	return fmt.Sprintf("(%s) %s", fmtSize(b.Size(), b.GetUnit()), b.Name())
 }
 
 //////////////////////////////

dynamic_tree_test.go

@@ -3,74 +3,69 @@ package main
 import (
 	"fmt"
 	"testing"
+	"time"
 )
 
-const result = `(0) /
+const result = `(0.00 KB) /
- ├──(0) i1
+ ├──(0.00 KB) i1
- │   ├──(0) i1s1
+ │   ├──(0.00 KB) i1s1
- │   │   ├──(0) i1s1b1
+ │   │   ├──(0.00 KB) i1s1b1
- │   │   ├──(0) i1s1b2
+ │   │   ├──(0.00 KB) i1s1b2
- │   │   └──(0) i1s1b3
+ │   │   └──(0.00 KB) i1s1b3
- │   │      (0) i1s1b3
+ │   ├──(0.00 KB) i1s2
- │   ├──(0) i1s2
+ │   │   ├──(0.00 KB) i1s2b1
- │   │   ├──(0) i1s2b1
+ │   │   └──(0.00 KB) i1s2b2
- │   │   └──(0) i1s2b2
+ │   ├──(0.00 KB) i1s3
- │   │      (0) i1s2b2
+ │   │   ├──(0.00 KB) i1s3s1
- │   ├──(0) i1s3
+ │   │   │   ├──(0.00 KB) i1s3s1b1
- │   │   ├──(0) i1s3s1
+ │   │   │   ├──(0.00 KB) i1s3s1b2
- │   │   │   ├──(0) i1s3s1b1
+ │   │   │   └──(0.00 KB) i1s3s1b3
- │   │   │   ├──(0) i1s3s1b2
+ │   │   └──(0.00 KB) i1s3b1
- │   │   │   └──(0) i1s3s1b3
+ │   └──(0.00 KB) i1s4
- │   │   │      (0) i1s3s1b3
+ │       └──(0.00 KB) i1s4b1
- │   │   └──(0) i1s3b1
+ ├──(0.00 KB) i2
- │   │      (0) i1s3b1
+ │   ├──(0.00 KB) i2b1
- │   └──(0) i1s4
+ │   └──(0.00 KB) i2b2
- │            (0) i1s4b1
+ └──(0.00 KB) i3
- ├──(0) i2
+     ├──(0.00 KB) i3s1
- │   ├──(0) i2b1
+     │   └──(0.00 KB) i3s1b1
- │   └──(0) i2b2
+     ├──(0.00 KB) i3s2
- │      (0) i2b2
+     │   ├──(0.00 KB) i3s2b1
- └──(0) i3
+     │   └──(0.00 KB) i3s2b2
-       ├──(0) i3s1
+     └──(0.00 KB) i3s3
-       │   └──(0) i3s1b1
+         └──(0.00 KB) i3s3b1
-       │      (0) i3s1b1
-       ├──(0) i3s2
-       │   ├──(0) i3s2b1
-       │   └──(0) i3s2b2
-       │      (0) i3s2b2
-       └──(0) i3s3
-                (0) i3s3b1
 `
 
 func createTree() *Top {
 	t := NewTop("/")
-	i1 := t.newChild("i1", IntermediateType).(*Intermediate)
+	t.SetUnit("KB")
-	i2 := t.newChild("i2", IntermediateType).(*Intermediate)
+	i1 := t.newIntermediate("i1")
-	i3 := t.newChild("i3", IntermediateType).(*Intermediate)
+	i2 := t.newIntermediate("i2")
-	i1s1 := i1.newChild("i1s1", IntermediateType).(*Intermediate)
+	i3 := t.newIntermediate("i3")
-	i1s2 := i1.newChild("i1s2", IntermediateType).(*Intermediate)
+	i1s1 := i1.newIntermediate("i1s1")
-	i1s3 := i1.newChild("i1s3", IntermediateType).(*Intermediate)
+	i1s2 := i1.newIntermediate("i1s2")
-	i1s4 := i1.newChild("i1s4", IntermediateType).(*Intermediate)
+	i1s3 := i1.newIntermediate("i1s3")
-	i3s1 := i3.newChild("i3s1", IntermediateType).(*Intermediate)
+	i1s4 := i1.newIntermediate("i1s4")
-	i3s2 := i3.newChild("i3s2", IntermediateType).(*Intermediate)
+	i3s1 := i3.newIntermediate("i3s1")
-	i3s3 := i3.newChild("i3s3", IntermediateType).(*Intermediate)
+	i3s2 := i3.newIntermediate("i3s2")
-	_ = i1s1.newChild("i1s1b1", BottomType).(*Bottom)
+	i3s3 := i3.newIntermediate("i3s3")
-	_ = i1s1.newChild("i1s1b2", BottomType).(*Bottom)
+	_ = i1s1.newBottom("i1s1b1")
-	_ = i1s1.newChild("i1s1b3", BottomType).(*Bottom)
+	_ = i1s1.newBottom("i1s1b2")
-	_ = i1s2.newChild("i1s2b1", BottomType).(*Bottom)
+	_ = i1s1.newBottom("i1s1b3")
-	_ = i1s2.newChild("i1s2b2", BottomType).(*Bottom)
+	_ = i1s2.newBottom("i1s2b1")
-	i1s3s1 := i1s3.newChild("i1s3s1", IntermediateType).(*Intermediate)
+	_ = i1s2.newBottom("i1s2b2")
-	_ = i1s3.newChild("i1s3b1", BottomType).(*Bottom)
+	i1s3s1 := i1s3.newIntermediate("i1s3s1")
-	_ = i1s4.newChild("i1s4b1", BottomType).(*Bottom)
+	_ = i1s3.newBottom("i1s3b1")
-	_ = i2.newChild("i2b1", BottomType).(*Bottom)
+	_ = i1s4.newBottom("i1s4b1")
-	_ = i2.newChild("i2b2", BottomType).(*Bottom)
+	_ = i2.newBottom("i2b1")
-	_ = i3s1.newChild("i3s1b1", BottomType).(*Bottom)
+	_ = i2.newBottom("i2b2")
-	_ = i3s2.newChild("i3s2b1", BottomType).(*Bottom)
+	_ = i3s1.newBottom("i3s1b1")
-	_ = i3s2.newChild("i3s2b2", BottomType).(*Bottom)
+	_ = i3s2.newBottom("i3s2b1")
-	_ = i3s3.newChild("i3s3b1", BottomType).(*Bottom)
+	_ = i3s2.newBottom("i3s2b2")
-	_ = i1s3s1.newChild("i1s3s1b1", BottomType).(*Bottom)
+	_ = i3s3.newBottom("i3s3b1")
-	_ = i1s3s1.newChild("i1s3s1b2", BottomType).(*Bottom)
+	_ = i1s3s1.newBottom("i1s3s1b1")
-	_ = i1s3s1.newChild("i1s3s1b3", BottomType).(*Bottom)
+	_ = i1s3s1.newBottom("i1s3s1b2")
+	_ = i1s3s1.newBottom("i1s3s1b3")
 
 	return t
 }
@@ -79,6 +74,31 @@ func TestString(t *testing.T) {
 	tree := createTree()
 	repr := fmt.Sprintf("%s", tree)
 	if repr != result {
-		t.Errorf("%s", repr)
+		t.Errorf("repr:\n%s\n\nresult:\n%s", repr, result)
+	}
+}
+
+func TestSize(t *testing.T) {
+	root := NewTop(".")
+	i := NewIntermediate("i", root)
+	b := NewBottom("b", i)
+	go func() {
+		ch := b.Collector()
+		ch <- 1024
+	}()
+	go b.Collect()
+	go i.Collect()
+	go root.Collect()
+
+	time.Sleep(time.Second)
+
+	if bsize := b.Size(); bsize != 1024 {
+		t.Errorf("b -> wrong size: %d\n", bsize)
+	}
+	if isize := i.Size(); isize != 1024 {
+		t.Errorf("i -> wrong size: %d\n", isize)
+	}
+	if tsize := root.Size(); tsize != 1024 {
+		t.Errorf("t -> wrong size: %d\n", tsize)
 	}
 }

main.go

@@ -6,14 +6,50 @@ import (
 	"time"
 )
 
+type UnitValue struct {
+	unit string
+}
+
+func (u *UnitValue) String() string {
+	return u.unit
+}
+
+func (u *UnitValue) Set(value string) error {
+	switch value {
+	case "B":
+		u.unit = "B"
+		return nil
+	case "KB":
+		u.unit = "KB"
+		return nil
+	case "MB":
+		u.unit = "MB"
+		return nil
+	case "GB":
+		u.unit = "GB"
+		return nil
+	case "TB":
+		u.unit = "TB"
+		return nil
+	case "PB":
+		u.unit = "PB"
+		return nil
+	default:
+		return ErrUnknownUnit
+	}
+}
+
 func main() {
 	var path string
 	var depth int
+	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()
 
 	t := NewTop(path)
+	t.SetUnit(unit.String())
 	go t.Spawn(depth)
 	go t.Collect()
 	for {