-
Notifications
You must be signed in to change notification settings - Fork 7
/
cliflag.go
183 lines (156 loc) · 4.93 KB
/
cliflag.go
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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
// Package cliflag extends flagset with environment variable defaults.
//
// Usage:
//
// cliflag.String(root.Flags(), &address, "address", "a", "CODER_ADDRESS", "127.0.0.1:3000", "The address to serve the API and dashboard")
//
// Will produce the following usage docs:
//
// -a, --address string The address to serve the API and dashboard (uses $CODER_ADDRESS). (default "127.0.0.1:3000")
package cliflag
import (
"fmt"
"os"
"strconv"
"strings"
"time"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
)
// IsSetBool returns the value of the boolean flag if it is set.
// It returns false if the flag isn't set or if any error occurs attempting
// to parse the value of the flag.
func IsSetBool(cmd *cobra.Command, name string) bool {
val, ok := IsSet(cmd, name)
if !ok {
return false
}
b, err := strconv.ParseBool(val)
return err == nil && b
}
// IsSet returns the string value of the flag and whether it was set.
func IsSet(cmd *cobra.Command, name string) (string, bool) {
flag := cmd.Flag(name)
if flag == nil {
return "", false
}
return flag.Value.String(), flag.Changed
}
// String sets a string flag on the given flag set.
func String(flagset *pflag.FlagSet, name, shorthand, env, def, usage string) {
v, ok := os.LookupEnv(env)
if !ok || v == "" {
v = def
}
flagset.StringP(name, shorthand, v, fmtUsage(usage, env))
}
// StringVarP sets a string flag on the given flag set.
func StringVarP(flagset *pflag.FlagSet, p *string, name string, shorthand string, env string, def string, usage string) {
v, ok := os.LookupEnv(env)
if !ok || v == "" {
v = def
}
flagset.StringVarP(p, name, shorthand, v, fmtUsage(usage, env))
}
func StringArray(flagset *pflag.FlagSet, name, shorthand, env string, def []string, usage string) {
v, ok := os.LookupEnv(env)
if !ok || v == "" {
if v == "" {
def = []string{}
} else {
def = strings.Split(v, ",")
}
}
flagset.StringArrayP(name, shorthand, def, fmtUsage(usage, env))
}
func StringArrayVarP(flagset *pflag.FlagSet, ptr *[]string, name string, shorthand string, env string, def []string, usage string) {
val, ok := os.LookupEnv(env)
if ok {
if val == "" {
def = []string{}
} else {
def = strings.Split(val, ",")
}
}
flagset.StringArrayVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
}
// Uint8VarP sets a uint8 flag on the given flag set.
func Uint8VarP(flagset *pflag.FlagSet, ptr *uint8, name string, shorthand string, env string, def uint8, usage string) {
val, ok := os.LookupEnv(env)
if !ok || val == "" {
flagset.Uint8VarP(ptr, name, shorthand, def, fmtUsage(usage, env))
return
}
vi64, err := strconv.ParseUint(val, 10, 8)
if err != nil {
flagset.Uint8VarP(ptr, name, shorthand, def, fmtUsage(usage, env))
return
}
flagset.Uint8VarP(ptr, name, shorthand, uint8(vi64), fmtUsage(usage, env))
}
// IntVarP sets a uint8 flag on the given flag set.
func IntVarP(flagset *pflag.FlagSet, ptr *int, name string, shorthand string, env string, def int, usage string) {
val, ok := os.LookupEnv(env)
if !ok || val == "" {
flagset.IntVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
return
}
vi64, err := strconv.Atoi(val)
if err != nil {
flagset.IntVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
return
}
flagset.IntVarP(ptr, name, shorthand, vi64, fmtUsage(usage, env))
}
func Bool(flagset *pflag.FlagSet, name, shorthand, env string, def bool, usage string) {
val, ok := os.LookupEnv(env)
if !ok || val == "" {
flagset.BoolP(name, shorthand, def, fmtUsage(usage, env))
return
}
valb, err := strconv.ParseBool(val)
if err != nil {
flagset.BoolP(name, shorthand, def, fmtUsage(usage, env))
return
}
flagset.BoolP(name, shorthand, valb, fmtUsage(usage, env))
}
// BoolVarP sets a bool flag on the given flag set.
func BoolVarP(flagset *pflag.FlagSet, ptr *bool, name string, shorthand string, env string, def bool, usage string) {
val, ok := os.LookupEnv(env)
if !ok || val == "" {
flagset.BoolVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
return
}
valb, err := strconv.ParseBool(val)
if err != nil {
flagset.BoolVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
return
}
flagset.BoolVarP(ptr, name, shorthand, valb, fmtUsage(usage, env))
}
// DurationVarP sets a time.Duration flag on the given flag set.
func DurationVarP(flagset *pflag.FlagSet, ptr *time.Duration, name string, shorthand string, env string, def time.Duration, usage string) {
val, ok := os.LookupEnv(env)
if !ok || val == "" {
flagset.DurationVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
return
}
valb, err := time.ParseDuration(val)
if err != nil {
flagset.DurationVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
return
}
flagset.DurationVarP(ptr, name, shorthand, valb, fmtUsage(usage, env))
}
func fmtUsage(u string, env string) string {
if env != "" {
// Avoid double dotting.
dot := "."
if strings.HasSuffix(u, ".") {
dot = ""
}
u = fmt.Sprintf("%s%s\n(Consumes $%s", u, dot, env)
}
return u
}