Filter by size limits

cmd/ru/main.go

@@ -44,6 +44,30 @@ func (u *UnitValue) Set(value string) error {
 	}
 }
 
+type LimitValue struct {
+	limit int64
+	input string
+}
+
+func (l *LimitValue) String() string {
+	return fmt.Sprintf("%d", l.limit)
+}
+
+func (l *LimitValue) Set(size string) error {
+	value, err := tree.ParseSize(size)
+	l.limit = value
+	l.input = size
+	return err
+}
+
+func validateLimits(min, max *LimitValue) bool {
+	if max.limit != 0 && min.limit > max.limit {
+		fmt.Fprintf(os.Stderr, "Given min-limit (%s) is above max-limit (%s)\n", min.input, max.input)
+		return false
+	}
+	return true
+}
+
 func min(a, b int) int {
 	if a < b {
 		return a
@@ -57,9 +81,13 @@ func main() {
 	var root tree.Node
 	var interval time.Duration
 	var unit = &UnitValue{unit: "KB"}
+	var minLimit = &LimitValue{limit: 0}
+	var maxLimit = &LimitValue{limit: 0}
 	cli := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
 	cli.IntVar(&depth, "depth", 0, "Depth to display")
 	cli.Var(unit, "unit", "Unit in which to report size")
+	cli.Var(minLimit, "limit-min", "Mininmum size of directories to display, with unit (0 means no limit)")
+	cli.Var(maxLimit, "limit-max", "Maximum size of directories to display, with unit (0 means no limit)")
 	cli.DurationVar(&interval, "interval", 100*time.Millisecond, "The update interval")
 	cli.Usage = func() {
 		fmt.Fprintf(cli.Output(), "Usage:\n%s [opts] [PATH]\n\n  PATH: the root path to start from. Defaults to $PWD.\n\nopts:\n", os.Args[0])
@@ -67,6 +95,10 @@ func main() {
 	}
 	cli.Parse(os.Args[1:])
 
+	if !validateLimits(minLimit, maxLimit) {
+		os.Exit(-1)
+	}
+
 	switch narg := cli.NArg(); narg {
 	case 0:
 		path = "."
@@ -83,6 +115,7 @@ func main() {
 		root = tree.NewTop(path)
 	}
 	root.SetUnit(unit.String())
+	root.SetLimits(minLimit.limit, maxLimit.limit)
 	out := log.New(os.Stdout, "", 0)
 
 	go root.Spawn(depth)
@@ -92,18 +125,24 @@ func main() {
 		os.Exit(1)
 	}
 
-	firstRound := true
+	treeDepth := 0
 	for {
 		select {
 		case <-time.After(interval):
-			if !firstRound {
+			if treeDepth != 0 {
-				depth := min(root.Depth()+2, height)
+				out.Printf("\r\033[%dA", treeDepth)
-				out.Printf("\r\033[%dA", depth)
 			} else {
 				out.Print("\r\033[A")
 			}
+			newTreeDepth := min(root.Depth()+2, height)
 			out.Print(root)
-			firstRound = false
+			if newTreeDepth <= treeDepth {
+				//for i := treeDepth; i < newTreeDepth; i++ {
+				//	out.Print("\n\033[2K")
+				//}
+				out.Print("\033[J\033[A")
+			}
+			treeDepth = newTreeDepth
 			if root.Complete() {
 				return
 			}

tree/dynamic_tree.go

@@ -19,10 +19,11 @@ var ErrShouldBeBottom = errors.New("should be of type Bottom")
 var ErrChannelError = errors.New("channel errored")
 var ErrUnknownUnit = errors.New("unknown unit")
 var (
-	NodeB   = "│"
+	NodeB     = "│"
-	NodeT   = "├──"
+	NodeT     = "├──"
-	NodeL   = "└──"
+	NodeL     = "└──"
-	NodePad = "   "
+	NodePad   = "   "
+	CloseLine = "\033[K\n"
 )
 
 type Node interface {
@@ -32,6 +33,7 @@ type Node interface {
 	Spawn(int) error
 	Collect() error
 	AddCollector(chan int64)
+	SetLimits(int64, int64)
 	Level() int
 	Name() string
 	Complete() bool
@@ -48,6 +50,8 @@ type Top struct {
 	unit     ByteSize
 	tree     []Node
 	collect  []<-chan int64
+	minLimit int64
+	maxLimit int64
 	level    int
 	complete bool
 }
@@ -58,7 +62,7 @@ func NewTop(path string) *Top {
 }
 
 func (t *Top) SetUnit(unit string) {
-	t.unit = setUnit(unit)
+	t.unit = parseUnit(unit)
 }
 
 func (t *Top) GetUnit() ByteSize {
@@ -112,6 +116,11 @@ func (t *Top) AddCollector(collect chan int64) {
 	t.collect = append(t.collect, collect)
 }
 
+func (t *Top) SetLimits(min, max int64) {
+	t.minLimit = min
+	t.maxLimit = max
+}
+
 func (t *Top) Level() int {
 	return t.level
 }
@@ -132,7 +141,9 @@ func (t *Top) Complete() bool {
 func (t *Top) Depth() int {
 	var depth int
 	for _, child := range t.tree {
-		depth += child.Depth()
+		if isIntoLimits(child.Size(), t.minLimit, t.maxLimit) {
+			depth += child.Depth()
+		}
 	}
 	return depth
 }
@@ -140,22 +151,25 @@ func (t *Top) Depth() int {
 func (t *Top) String() string {
 	var out string
 	var lines []string
-	out += fmt.Sprintf("(%s) %s\n", fmtSize(t.Size(), t.GetUnit()), t.path)
+	out += fmt.Sprintf("(%s) %s%s", fmtSize(t.Size(), t.GetUnit()), t.path, CloseLine)
-	treeSize := len(t.tree)
+	treeSize := 0
 	for _, child := range t.tree {
-		lines = append(lines, fmt.Sprintf("%s", child))
+		if isIntoLimits(child.Size(), t.minLimit, t.maxLimit) {
+			lines = append(lines, fmt.Sprintf("%s", child))
+			treeSize += 1
+		}
 	}
 	for a := 0; a < treeSize-1; a++ {
 		childLines := strings.Split(lines[a], "\n")
-		out += fmt.Sprintf(" %s%s\n", NodeT, childLines[0])
+		out += fmt.Sprintf(" %s%s%s", NodeT, childLines[0], CloseLine)
 		for b := 1; b < len(childLines); b++ {
-			out += fmt.Sprintf(" %s%s\n", NodeB, childLines[b])
+			out += fmt.Sprintf(" %s%s%s", NodeB, childLines[b], CloseLine)
 		}
 	}
 	childLines := strings.Split(lines[treeSize-1], "\n")
-	out += fmt.Sprintf(" %s%s\n", NodeL, childLines[0])
+	out += fmt.Sprintf(" %s%s%s", NodeL, childLines[0], CloseLine)
 	for a := 1; a < len(childLines); a++ {
-		out += fmt.Sprintf("%s%s\n", "  ", childLines[a])
+		out += fmt.Sprintf("%s%s%s", "  ", childLines[a], CloseLine)
 	}
 	return out
 }
@@ -201,6 +215,8 @@ type Intermediate struct {
 	tree     []Node
 	collect  []<-chan int64
 	refer    chan int64
+	minLimit int64
+	maxLimit int64
 	level    int
 	complete bool
 }
@@ -215,7 +231,7 @@ func NewIntermediate(path string, parent Node) (*Intermediate, chan int64) {
 }
 
 func (i *Intermediate) SetUnit(unit string) {
-	i.unit = setUnit(unit)
+	i.unit = parseUnit(unit)
 }
 
 func (i *Intermediate) GetUnit() ByteSize {
@@ -272,6 +288,11 @@ func (i *Intermediate) AddCollector(collect chan int64) {
 	i.collect = append(i.collect, collect)
 }
 
+func (i *Intermediate) SetLimits(min, max int64) {
+	i.minLimit = min
+	i.maxLimit = max
+}
+
 func (i *Intermediate) Level() int {
 	return i.level
 }
@@ -292,7 +313,9 @@ func (i *Intermediate) Complete() bool {
 func (i *Intermediate) Depth() int {
 	var depth int
 	for _, child := range i.tree {
-		depth += child.Depth()
+		if isIntoLimits(child.Size(), i.minLimit, i.maxLimit) {
+			depth += child.Depth()
+		}
 	}
 	return depth
 }
@@ -300,25 +323,28 @@ func (i *Intermediate) Depth() int {
 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)
+	treeSize := 0
 	for _, child := range i.tree {
-		lines = append(lines, fmt.Sprintf("%s", child))
+		if isIntoLimits(child.Size(), i.minLimit, i.maxLimit) {
+			lines = append(lines, fmt.Sprintf("%s", child))
+			treeSize += 1
+		}
 	}
 
 	for a := 0; a < treeSize-1; a++ {
 		childLines := strings.Split(lines[a], "\n")
-		out += fmt.Sprintf("%s%s%s\n", NodePad, NodeT, childLines[0])
+		out += fmt.Sprintf("%s%s%s", NodePad, NodeT, childLines[0])
 		for b := 1; b < len(childLines); b++ {
-			out += fmt.Sprintf("%s%s%s\n", NodePad, NodeB, childLines[b])
+			out += fmt.Sprintf("%s%s%s", 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])
+		out += fmt.Sprintf("%s%s%s", 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[a])
 		}
 		out += fmt.Sprintf(" %s%s", NodePad, childLines[lenChildLines-1])
 	} else {
@@ -368,6 +394,8 @@ type Bottom struct {
 	walker   Walker
 	collect  chan int64
 	refer    chan int64
+	minLimit int64
+	maxLimit int64
 	level    int
 	complete bool
 }
@@ -381,7 +409,7 @@ func NewBottom(path string, parent Node) (*Bottom, chan int64) {
 }
 
 func (b *Bottom) SetUnit(unit string) {
-	b.unit = setUnit(unit)
+	b.unit = parseUnit(unit)
 }
 
 func (b *Bottom) GetUnit() ByteSize {
@@ -414,6 +442,11 @@ func (b *Bottom) AddCollector(collect chan int64) {
 	b.collect = collect
 }
 
+func (b *Bottom) SetLimits(min, max int64) {
+	b.minLimit = min
+	b.maxLimit = max
+}
+
 func (b *Bottom) Level() int {
 	return b.level
 }
@@ -455,7 +488,7 @@ func NewSingle(path string) *Single {
 }
 
 func (s *Single) SetUnit(unit string) {
-	s.unit = setUnit(unit)
+	s.unit = parseUnit(unit)
 }
 
 func (s *Single) GetUnit() ByteSize {
@@ -485,6 +518,11 @@ func (s *Single) AddCollector(collect chan int64) {
 	s.collect = collect
 }
 
+func (s *Single) SetLimits(min, max int64) {
+	// just ignore limits
+	return
+}
+
 func (s *Single) Level() int {
 	return 0
 }

tree/tree_utils.go

@@ -1,10 +1,17 @@
 package tree
 
 import (
+	"errors"
 	"fmt"
+	"math"
+	"regexp"
+	"strconv"
 	"sync"
 )
 
+var ErrUnparsableSize = errors.New("unparsable size")
+var matchSize = regexp.MustCompile("^(\\d+(\\.\\d+)?)([KMGTP]?B?)$")
+
 type ByteSize uint
 
 const (
@@ -20,7 +27,7 @@ func isLastLevel(node Node, maxLevel int) bool {
 	return maxLevel-1 == node.Level()
 }
 
-func setUnit(unit string) ByteSize {
+func parseUnit(unit string) ByteSize {
 	switch unit {
 	case "B":
 		return B
@@ -40,6 +47,19 @@ func setUnit(unit string) ByteSize {
 	}
 }
 
+func parseUnitSafe(unit string) (ByteSize, error) {
+	var returnVal ByteSize
+	var err error
+	defer func() {
+		if r := recover(); r != nil {
+			returnVal = 0
+			err = ErrUnparsableSize
+		}
+	}()
+	returnVal = parseUnit(unit)
+	return returnVal, err
+}
+
 func fmtUnit(unit ByteSize) string {
 	switch unit {
 	case B:
@@ -66,6 +86,36 @@ func fmtSize(size int64, unit ByteSize) string {
 	return fmt.Sprintf("%.2f %s", dimension, fmtUnit(unit))
 }
 
+func ParseSize(size string) (int64, error) {
+	result := matchSize.FindStringSubmatch(size)
+	sizeNum := result[1]
+	if sizeNum == "" {
+		return 0, ErrUnparsableSize
+	}
+	unit, err := parseUnitSafe(result[3])
+	if err != nil {
+		return 0, err
+	}
+	num, err := strconv.ParseFloat(sizeNum, 64)
+	if err != nil {
+		return 0, err
+	}
+
+	return int64(math.Floor(num * float64(unit))), nil
+}
+
+func isIntoLimits(size, min, max int64) bool {
+	if min <= size {
+		if max == 0 {
+			return true
+		}
+		if size <= max {
+			return true
+		}
+	}
+	return false
+}
+
 func merge(cs []<-chan int64) <-chan int64 {
 	var wg sync.WaitGroup
 	out := make(chan int64)