diff options
| -rw-r--r-- | flag.go | 49 | ||||
| -rw-r--r-- | flag_test.go | 59 |
2 files changed, 108 insertions, 0 deletions
@@ -55,6 +55,53 @@ func FlagTSliceVar[T any](fs *flag.FlagSet, p *[]T, name string, values []T, usa fs.Var(&flagValueSlice[T]{Parse: parse, Separator: sep, Values: p}, name, usage) } +// InitFlagSet initializes a flag.FlagSet by setting flags in the +// following order: environment variables, then an arbitrary map, then +// command line arguments. +// +// Note that InitFlagSet does not require the use of any of the Flag +// functions defined in this package. Standard flags will work just as +// well. +func InitFlagSet(fs *flag.FlagSet, env []string, cfg map[string]string, args []string) (err error) { + var environ map[string]string + if env != nil { + environ = make(map[string]string, len(env)) + for _, kv := range env { + if buf := strings.SplitN(kv, "=", 2); len(buf) == 2 { + environ[buf[0]] = buf[1] + continue + } + if val, ok := os.LookupEnv(kv); ok { + environ[kv] = val + } + } + } + + fs.VisitAll(func(f *flag.Flag) { + if err != nil { + return + } + + var next string + if val, found := environ[strings.ToUpper(strings.ReplaceAll(f.Name, "-", "_"))]; found { + next = val + } + if val, found := cfg[f.Name]; found { + next = val + } + if next != "" { + err = f.Value.Set(next) + } + if f, ok := f.Value.(interface{ resetShouldAppend() }); ok { + f.resetShouldAppend() + } + }) + if err == nil && !fs.Parsed() { + return fs.Parse(args) + } + return err +} + // ParseString returns the string passed with no error set. func ParseString(s string) (string, error) { return s, nil @@ -130,3 +177,5 @@ func (f *flagValueSlice[T]) String() string { } return fmt.Sprintf("%v", *f.Values) } + +func (f *flagValueSlice[T]) resetShouldAppend() { f.shouldAppend = false } diff --git a/flag_test.go b/flag_test.go index ccbfadb..584a2c9 100644 --- a/flag_test.go +++ b/flag_test.go @@ -52,3 +52,62 @@ func TestFlagTSliceVar(s *testing.T) { t.AssertErrorIs(nil, fs.Parse([]string{"-test=1", "-test=2", "-test=42,84"})) t.AssertEqual([]int{1, 2, 42, 84}, fl) } + +func TestInitFlagSet(s *testing.T) { + t := core.T{T: s} + + for _, tc := range []struct { + name string + env []string + cfg map[string]string + args []string + + expInt int + expIntSlice []int + expStr string + expErr error + }{ + { + name: "ArgsOnly", + args: []string{"-int=42", "-int-slice=42,84"}, + + expInt: 42, + expIntSlice: []int{42, 84}, + }, + { + name: "EnvOnly", + env: []string{"INT=42", "INT_SLICE=42,84"}, + + expInt: 42, + expIntSlice: []int{42, 84}, + }, + { + name: "CfgOnly", + cfg: map[string]string{"int": "42", "int-slice": "42,84"}, + + expInt: 42, + expIntSlice: []int{42, 84}, + }, + { + name: "InOrder", + env: []string{"STRING=Hello World!"}, + cfg: map[string]string{"string": "Hello Universe!", "int-slice": "42,84"}, + args: []string{"-int=42", "-int-slice=21,42"}, + + expInt: 42, + expIntSlice: []int{21, 42}, + expStr: "Hello Universe!", + }, + } { + t.Run(tc.name, func(t *core.T) { + fs := flag.NewFlagSet("", flag.PanicOnError) + fi := fs.Int("int", 0, "") + fl := core.FlagTSlice(fs, "int-slice", nil, "", strconv.Atoi, ",") + fm := fs.String("string", "", "") + t.AssertErrorIs(tc.expErr, core.InitFlagSet(fs, tc.env, tc.cfg, tc.args)) + t.AssertEqual(tc.expInt, *fi) + t.AssertEqual(tc.expIntSlice, *fl) + t.AssertEqual(tc.expStr, *fm) + }) + } +} |