ruspa/cmd/ru/main.go

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"time"

	"git.lattuga.net/blallo/ruspa/tree"
	"golang.org/x/crypto/ssh/terminal"
)

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 tree.ErrUnknownUnit
	}
}

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
	}
	return b
}

func main() {
	var path string
	var depth int
	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])
		cli.PrintDefaults()
	}
	cli.Parse(os.Args[1:])

	if !validateLimits(minLimit, maxLimit) {
		os.Exit(-1)
	}

	switch narg := cli.NArg(); narg {
	case 0:
		path = "."
	case 1:
		path = cli.Args()[0]
	default:
		fmt.Fprintln(os.Stderr, "Too many arguments")
		os.Exit(-1)
	}

	if depth == 0 || !tree.AnyDirectoryDownThere(path) {
		root = tree.NewSingle(path)
	} else {
		root = tree.NewTop(path)
	}
	root.SetUnit(unit.String())
	root.SetLimits(minLimit.limit, maxLimit.limit)
	out := log.New(os.Stdout, "", 0)

	go root.Spawn(depth)
	_, height, err := terminal.GetSize(int(os.Stdin.Fd()))
	if err != nil {
		fmt.Fprintln(os.Stderr, "Could not get terminal size")
		os.Exit(1)
	}

	treeDepth := 0
	for {
		select {
		case <-time.After(interval):
			if treeDepth != 0 {
				out.Printf("\r\033[%dA", treeDepth)
			} else {
				out.Print("\r\033[A")
			}
			newTreeDepth := min(root.Depth()+2, height)
			out.Print(root)
			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
			}
		}
	}
}