Last active
June 14, 2019 21:15
-
-
Save nilium/b49e2d1ace71b658564a38dca3303f19 to your computer and use it in GitHub Desktop.
Example of parsing environment variables and CLI flags using the flag package
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package main | |
import ( | |
"flag" | |
"log" | |
"os" | |
"strings" | |
"unicode" | |
) | |
func main() { | |
// f communicates program name and default log output target (stderr in this case). | |
f := flag.NewFlagSet("buildit", flag.ContinueOnError) | |
// Bind flags for configuration | |
id := f.Int("backend-id", 100, "Backend agent `ID`.") | |
// Parse CLI flags | |
if err := f.Parse(os.Args[1:]); err == flag.ErrHelp { | |
os.Exit(2) | |
} else if err != nil { | |
log.Printf("Error parsing CLI config: %v", err) | |
os.Exit(1) | |
} | |
// Load environment variables for unset CLI flags | |
if err := LoadEnvironmentConfig(f, os.LookupEnv); err != nil { | |
log.Printf("Error parsing environment config: %v", err) | |
os.Exit(1) | |
} | |
log.Printf("ID = %v", *id) // DEBUG | |
} | |
// EnvFunc is an os.LookupEnv-compatible function type. | |
// An implementation looks up a value by name and return both the value (if any) and whether the | |
// name was known. | |
type EnvFunc func(name string) (value string, ok bool) | |
// LoadEnvironmentConfig assigns environment variables as values for unset flags defined in f. | |
// Flags are set in lexical order. | |
// | |
// If lookup is nil, os.LookupEnv is used instead. | |
// | |
// If any errors occur in setting a flag, the first error is returned. | |
func LoadEnvironmentConfig(f *flag.FlagSet, lookup EnvFunc) error { | |
if lookup == nil { | |
lookup = os.LookupEnv | |
} | |
known := map[string]struct{}{} | |
f.Visit(func(v *flag.Flag) { | |
known[v.Name] = struct{}{} | |
}) | |
var err error | |
prefix := f.Name() + "_" | |
f.VisitAll(func(v *flag.Flag) { | |
if _, ok := known[v.Name]; ok { | |
return | |
} | |
name := envVarName(prefix + v.Name) | |
val, ok := lookup(name) | |
if !ok { | |
return | |
} | |
// Take only the first error | |
if ferr := f.Set(v.Name, val); err == nil && ferr != nil { | |
err = ferr | |
} | |
}) | |
return err | |
} | |
// envVarName accepts a string and returns an enviornment variable name of the form NAME_FOO. | |
// | |
// The resulting name does not contain leading underscores or numbers. Repeated underscores are | |
// trimmed. The name may end in an underscore. | |
// | |
// Characters that are not in the range [A-Za-z0-9_] are replaced with underscores. | |
func envVarName(name string) string { | |
var last rune = -1 | |
return strings.Map(func(r rune) rune { | |
r = unicode.ToUpper(r) | |
alpha := r >= 'A' && r <= 'Z' | |
numeric := r >= '0' && r <= '9' | |
if !(alpha || numeric) { | |
r = '_' | |
} | |
if (last == -1 && !alpha) || (last == '_' && r == '_') { | |
return -1 | |
} | |
last = r | |
return r | |
}, name) | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment