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

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"}
	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.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:])

	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())
	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)
	}

	firstRound := true
	for {
		select {
		case <-time.After(interval):
			if !firstRound {
				depth := min(root.Depth()+2, height)
				out.Printf("\r\033[%dA", depth)
			} else {
				out.Print("\r\033[A")
			}
			out.Print(root)
			firstRound = false
			if root.Complete() {
				return
			}
		}
	}
}