1 package xmain
2
3 import (
4 "bytes"
5 "fmt"
6 "io"
7 "strconv"
8 "strings"
9
10 "github.com/spf13/pflag"
11
12 "oss.terrastruct.com/util-go/xos"
13 )
14
15 type Opts struct {
16 Args []string
17 Flags *pflag.FlagSet
18 env *xos.Env
19
20 flagEnv map[string]string
21 }
22
23 func NewOpts(env *xos.Env, args []string) *Opts {
24 flags := pflag.NewFlagSet("", pflag.ContinueOnError)
25 flags.SortFlags = false
26 flags.Usage = func() {}
27 flags.SetOutput(io.Discard)
28 return &Opts{
29 Args: args,
30 Flags: flags,
31 env: env,
32 flagEnv: make(map[string]string),
33 }
34 }
35
36
37
38 func (o *Opts) Defaults() string {
39 buf := new(bytes.Buffer)
40
41 var lines []string
42
43 maxlen := 0
44 maxEnvLen := 0
45 o.Flags.VisitAll(func(flag *pflag.Flag) {
46 if flag.Hidden {
47 return
48 }
49
50 line := ""
51 if flag.Shorthand != "" && flag.ShorthandDeprecated == "" {
52 line = fmt.Sprintf(" -%s, --%s", flag.Shorthand, flag.Name)
53 } else {
54 line = fmt.Sprintf(" --%s", flag.Name)
55 }
56
57 varname, usage := pflag.UnquoteUsage(flag)
58 if varname != "" {
59 line += " " + varname
60 }
61 if flag.NoOptDefVal != "" {
62 switch flag.Value.Type() {
63 case "string":
64 line += fmt.Sprintf("[=\"%s\"]", flag.NoOptDefVal)
65 case "bool":
66 if flag.NoOptDefVal != "true" {
67 line += fmt.Sprintf("[=%s]", flag.NoOptDefVal)
68 }
69 case "count":
70 if flag.NoOptDefVal != "+1" {
71 line += fmt.Sprintf("[=%s]", flag.NoOptDefVal)
72 }
73 default:
74 line += fmt.Sprintf("[=%s]", flag.NoOptDefVal)
75 }
76 }
77
78 line += "\x00"
79
80 if len(line) > maxlen {
81 maxlen = len(line)
82 }
83
84 if e, ok := o.flagEnv[flag.Name]; ok {
85 line += fmt.Sprintf("$%s", e)
86 }
87
88 line += "\x01"
89
90 if len(line) > maxEnvLen {
91 maxEnvLen = len(line)
92 }
93
94 line += usage
95 if flag.Value.Type() == "string" {
96 line += fmt.Sprintf(" (default %q)", flag.DefValue)
97 } else {
98 line += fmt.Sprintf(" (default %s)", flag.DefValue)
99 }
100 if len(flag.Deprecated) != 0 {
101 line += fmt.Sprintf(" (DEPRECATED: %s)", flag.Deprecated)
102 }
103
104 lines = append(lines, line)
105 })
106
107 for _, line := range lines {
108 sidx1 := strings.Index(line, "\x00")
109 sidx2 := strings.Index(line, "\x01")
110 spacing1 := strings.Repeat(" ", maxlen-sidx1)
111 spacing2 := strings.Repeat(" ", (maxEnvLen-maxlen)-sidx2+sidx1)
112 fmt.Fprintln(buf, line[:sidx1], spacing1, line[sidx1+1:sidx2], spacing2, wrap(maxEnvLen+3, 0, line[sidx2+1:]))
113 }
114
115 return buf.String()
116 }
117
118 func (o *Opts) getEnv(flag, k string) string {
119 if k != "" {
120 o.flagEnv[flag] = k
121 return o.env.Getenv(k)
122 }
123 return ""
124 }
125
126 func (o *Opts) Int64(envKey, flag, shortFlag string, defaultVal int64, usage string) (*int64, error) {
127 if env := o.getEnv(flag, envKey); env != "" {
128 envVal, err := strconv.ParseInt(env, 10, 64)
129 if err != nil {
130 return nil, UsageErrorf(`invalid environment variable %s. Expected int64. Found "%v".`, envKey, envVal)
131 }
132 defaultVal = envVal
133 }
134
135 return o.Flags.Int64P(flag, shortFlag, defaultVal, usage), nil
136 }
137
138 func (o *Opts) Int64Slice(envKey, flag, shortFlag string, defaultVal []int64, usage string) (*[]int64, error) {
139 if env := o.getEnv(flag, envKey); env != "" {
140 split := strings.Split(env, ",")
141 defaultVal = make([]int64, len(split))
142 for i, part := range split {
143 val, err := strconv.ParseInt(strings.TrimSpace(part), 10, 64)
144 if err != nil {
145 return nil, UsageErrorf(`invalid environment variable %s. Expected []int64. Found "%v".`, envKey, env)
146 }
147 defaultVal[i] = val
148 }
149 }
150
151 return o.Flags.Int64SliceP(flag, shortFlag, defaultVal, usage), nil
152 }
153
154 func (o *Opts) Float64(envKey, flag, shortFlag string, defaultVal float64, usage string) (*float64, error) {
155 if env := o.getEnv(flag, envKey); env != "" {
156 envVal, err := strconv.ParseFloat(env, 64)
157 if err != nil {
158 return nil, UsageErrorf(`invalid environment variable %s. Expected float64. Found "%v".`, envKey, envVal)
159 }
160 defaultVal = envVal
161 }
162
163 return o.Flags.Float64P(flag, shortFlag, defaultVal, usage), nil
164 }
165
166 func (o *Opts) String(envKey, flag, shortFlag string, defaultVal, usage string) *string {
167 if env := o.getEnv(flag, envKey); env != "" {
168 defaultVal = env
169 }
170
171 return o.Flags.StringP(flag, shortFlag, defaultVal, usage)
172 }
173
174 func (o *Opts) Bool(envKey, flag, shortFlag string, defaultVal bool, usage string) (*bool, error) {
175 if env := o.getEnv(flag, envKey); env != "" {
176 if !boolyEnv(env) {
177 return nil, UsageErrorf(`invalid environment variable %s. Expected bool. Found "%s".`, envKey, env)
178 }
179 if truthyEnv(env) {
180 defaultVal = true
181 } else {
182 defaultVal = false
183 }
184 }
185
186 return o.Flags.BoolP(flag, shortFlag, defaultVal, usage), nil
187 }
188
189 func boolyEnv(s string) bool {
190 return falseyEnv(s) || truthyEnv(s)
191 }
192
193 func falseyEnv(s string) bool {
194 return s == "0" || s == "false"
195 }
196
197 func truthyEnv(s string) bool {
198 return s == "1" || s == "true"
199 }
200
View as plain text