nmkill/flags.go at main · moaqz/nmkill · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
package main

import (
	"flag"
	"fmt"
	"strings"

	"github.com/docker/go-units"
)

const (
	bold      = "\033[1m"
	underline = "\033[4m"
	reset     = "\033[0m"
)

type Flags struct {
	Directory string
	Version   bool
	MinSize   string
	Help      bool
}

func (f *Flags) GetMinSize() (int64, error) {
	if f.MinSize == "" {
		return 0, nil
	}

	return units.FromHumanSize(f.MinSize)
}

func (f *Flags) Usage() {
	fmt.Print(customUsage())
}

func (f *Flags) define() {
	flag.StringVar(&f.Directory, "directory", ".", "Root directory to search (default: current directory)")
	flag.StringVar(&f.Directory, "d", ".", "Root directory to search (default: current directory)")

	flag.StringVar(&f.MinSize, "min-size", "", "Set minimum size")
	flag.StringVar(&f.MinSize, "s", "", "Set minimum size")

	flag.BoolVar(&f.Version, "v", false, "Print version")
	flag.BoolVar(&f.Version, "version", false, "Print version")

	flag.BoolVar(&f.Help, "h", false, "Print help")
	flag.BoolVar(&f.Help, "help", false, "Print help")
}

func customUsage() string {
	type flagGroup struct {
		names    []string
		defValue string
	}

	groups := make(map[string]*flagGroup)

	flag.VisitAll(func(f *flag.Flag) {
		if val, ok := groups[f.Usage]; !ok {
			groups[f.Usage] = &flagGroup{
				names:    []string{f.Name},
				defValue: f.DefValue,
			}
		} else {
			val.names = append(val.names, f.Name)
		}
	})

	var b strings.Builder

	fmt.Fprintf(&b, "%s%sUsage:%s %s [OPTIONS]\n\n", underline, bold, reset, "nmkill")
	fmt.Fprintf(&b, "%s%sOptions:%s\n", underline, bold, reset)

	// TODO: Preserve insertion order.
	// Currently map iteration is random, causing inconsistent help output.
	for usage, fgroup := range groups {
		fmt.Fprintf(&b, "  ")
		for i, name := range fgroup.names {
			if i > 0 {
				fmt.Fprintf(&b, ", ")
			}

			prefix := "-"
			if len(name) > 1 {
				prefix = "--"
			}

			fmt.Fprintf(&b, "%s%s", prefix, name)
		}

		fmt.Fprintf(&b, "\n\t%s", usage)
		b.WriteString("\n")
	}

	return b.String()
}

func ParseFlags() *Flags {
	f := Flags{}
	f.define()

	flag.Usage = func() {
		fmt.Print(customUsage())
	}

	flag.Parse()
	return &f
}