lib/envflag: small refactoring after 518c340ae3 and 02096e06d0

This commit is contained in:
Aliaksandr Valialkin 2022-10-29 02:28:58 +03:00
parent 76e8888272
commit 320ae1c60a
No known key found for this signature in database
GPG Key ID: A72BEC6CD3D0DED1

View File

@ -22,9 +22,46 @@ var (
//
// This function must be called instead of flag.Parse() before using any flags in the program.
func Parse() {
// Substitute %{ENV_VAR} inside args with the corresponding environment variable values
args := os.Args[1:]
dstArgs := args[:0]
ParseFlagSet(flag.CommandLine, os.Args[1:])
}
// ParseFlagSet parses the given args into the given fs.
func ParseFlagSet(fs *flag.FlagSet, args []string) {
args = expandArgs(args)
if err := fs.Parse(args); err != nil {
// Do not use lib/logger here, since it is uninitialized yet.
log.Fatalf("cannot parse flags %q: %s", args, err)
}
if !*enable {
return
}
// Remember explicitly set command-line flags.
flagsSet := make(map[string]bool)
fs.Visit(func(f *flag.Flag) {
flagsSet[f.Name] = true
})
// Obtain the remaining flag values from environment vars.
fs.VisitAll(func(f *flag.Flag) {
if flagsSet[f.Name] {
// The flag is explicitly set via command-line.
return
}
// Get flag value from environment var.
fname := getEnvFlagName(f.Name)
if v, ok := envtemplate.LookupEnv(fname); ok {
if err := fs.Set(f.Name, v); err != nil {
// Do not use lib/logger here, since it is uninitialized yet.
log.Fatalf("cannot set flag %s to %q, which is read from env var %q: %s", f.Name, v, fname, err)
}
}
})
}
// expandArgs substitutes %{ENV_VAR} placeholders inside args
// with the corresponding environment variable values.
func expandArgs(args []string) []string {
dstArgs := make([]string, 0, len(args))
for _, arg := range args {
s, err := envtemplate.ReplaceString(arg)
if err != nil {
@ -35,35 +72,7 @@ func Parse() {
dstArgs = append(dstArgs, s)
}
}
os.Args = os.Args[:1+len(dstArgs)]
// Parse flags
flag.Parse()
if !*enable {
return
}
// Remember explicitly set command-line flags.
flagsSet := make(map[string]bool)
flag.Visit(func(f *flag.Flag) {
flagsSet[f.Name] = true
})
// Obtain the remaining flag values from environment vars.
flag.VisitAll(func(f *flag.Flag) {
if flagsSet[f.Name] {
// The flag is explicitly set via command-line.
return
}
// Get flag value from environment var.
fname := getEnvFlagName(f.Name)
if v, ok := envtemplate.LookupEnv(fname); ok {
if err := flag.Set(f.Name, v); err != nil {
// Do not use lib/logger here, since it is uninitialized yet.
log.Fatalf("cannot set flag %s to %q, which is read from environment variable %q: %s", f.Name, v, fname, err)
}
}
})
return dstArgs
}
func getEnvFlagName(s string) string {