diff --git a/cmd/ru/main.go b/cmd/ru/main.go index 2b68eff..952c439 100644 --- a/cmd/ru/main.go +++ b/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 { - depth := min(root.Depth()+2, height) - out.Printf("\r\033[%dA", depth) + if treeDepth != 0 { + out.Printf("\r\033[%dA", treeDepth) } 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 } diff --git a/tree/dynamic_tree.go b/tree/dynamic_tree.go index 9dba6ca..48bd0d2 100644 --- a/tree/dynamic_tree.go +++ b/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 = "│" - NodeT = "├──" - NodeL = "└──" - NodePad = " " + NodeB = "│" + NodeT = "├──" + NodeL = "└──" + 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) - treeSize := len(t.tree) + out += fmt.Sprintf("(%s) %s%s", fmtSize(t.Size(), t.GetUnit()), t.path, CloseLine) + 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 } diff --git a/tree/tree_utils.go b/tree/tree_utils.go index 30e5d28..f85b38f 100644 --- a/tree/tree_utils.go +++ b/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)