envflag

Configure flags from the environment
Log | Files | Refs | README | LICENSE

flag.go (3519B)

      1 package envflag
      2 
      3 import (
      4 	"flag"
      5 	"fmt"
      6 	"os"
      7 	"regexp"
      8 	"strings"
      9 )
     10 
     11 // Option type for configing the parsing.
     12 type Option option
     13 type option func(*config)
     14 
     15 // FlagConverterFunc functions are used to update flag name for envvar lookups.
     16 type FlagConverterFunc func(flag string) string
     17 
     18 // UsageUpdaterFunc functions are used to update flag usage strings.
     19 type UsageUpdaterFunc func(key, usage string) string
     20 
     21 type config struct {
     22 	prefix        string
     23 	getenv        func(string) string
     24 	flagConverter FlagConverterFunc
     25 	usageUpdater  UsageUpdaterFunc
     26 }
     27 
     28 // FlagConverter converts the flag name to an environment variable key.
     29 func FlagConverter(f FlagConverterFunc) Option {
     30 	return func(cfg *config) {
     31 		cfg.flagConverter = f
     32 	}
     33 }
     34 
     35 // Prefix sets the prefix for all envvars.
     36 // This is applied after the FlagConverter function.
     37 func Prefix(s string) Option {
     38 	return func(cfg *config) {
     39 		cfg.prefix = s
     40 	}
     41 }
     42 
     43 // Getenv is used to fetch the environment, defaults to os.Getenv.
     44 func Getenv(f func(string) string) Option {
     45 	return func(cfg *config) {
     46 		cfg.getenv = f
     47 	}
     48 }
     49 
     50 // UsageUpdater enables configurable flag usage updates.
     51 func UsageUpdater(f UsageUpdaterFunc) Option {
     52 	return func(cfg *config) {
     53 		cfg.usageUpdater = f
     54 	}
     55 }
     56 
     57 // UsagePrefixer prefixes the flag usage with [envvar].
     58 //
     59 // Example output:
     60 //
     61 //	Usage of example:
     62 //	  -verbose
     63 //	        [VERBOSE] Be verbose
     64 func UsagePrefixer() Option {
     65 	return func(cfg *config) {
     66 		cfg.usageUpdater = usagePrefixer
     67 	}
     68 }
     69 
     70 // UsageSuffixer suffixes the flag usage with [envvar].
     71 //
     72 // Example output:
     73 //
     74 //	Usage of example:
     75 //	  -verbose
     76 //	        Be verbose [VERBOSE]
     77 func UsageSuffixer() Option {
     78 	return func(cfg *config) {
     79 		cfg.usageUpdater = usageSuffixer
     80 	}
     81 }
     82 
     83 // ParseWithEnv parses the command-line flags from os.Args[1:]. Must be called
     84 // after all flags are defined and before flags are accessed by the program.
     85 func Parse(opts ...Option) error {
     86 	return parseFlagSetWithEnv(flag.CommandLine, os.Args[1:], opts...)
     87 }
     88 
     89 // ParseFlagSetWithEnv parses flag definitions from the argument list, which
     90 // should not include the command name. Must be called after all flags in the
     91 // FlagSet are defined and before flags are accessed by the program. The return
     92 // value will be ErrHelp if -help or -h were set but not defined.
     93 func ParseFlagSet(fs *flag.FlagSet, arguments []string, opts ...Option) error {
     94 	return parseFlagSetWithEnv(fs, arguments, opts...)
     95 }
     96 
     97 func parseFlagSetWithEnv(fs *flag.FlagSet, arguments []string, opts ...Option) error {
     98 	if fs.Parsed() {
     99 		return fmt.Errorf("flag has already been parsed")
    100 	}
    101 	cfg := &config{
    102 		flagConverter: flagToEnv(),
    103 		getenv:        os.Getenv,
    104 	}
    105 	for _, o := range opts {
    106 		o(cfg)
    107 	}
    108 	var nerr error
    109 	fs.VisitAll(func(f *flag.Flag) {
    110 		envKey := cfg.flagConverter(f.Name)
    111 		envKey = cfg.prefix + envKey
    112 		if cfg.usageUpdater != nil {
    113 			f.Usage = cfg.usageUpdater(envKey, f.Usage)
    114 		}
    115 		if envVal := cfg.getenv(envKey); envVal != "" {
    116 			err := f.Value.Set(envVal)
    117 			if err != nil && nerr == nil {
    118 				nerr = err
    119 			}
    120 		}
    121 	})
    122 	if nerr != nil {
    123 		return nerr
    124 	}
    125 	return fs.Parse(arguments)
    126 }
    127 
    128 // convert this-format to THIS_FORMAT
    129 func flagToEnv() func(string) string {
    130 	re := regexp.MustCompile("[^a-zA-Z0-9_]+")
    131 	return func(f string) string {
    132 		return strings.ToUpper(re.ReplaceAllLiteralString(f, "_"))
    133 	}
    134 }
    135 
    136 func usageSuffixer(key, usage string) string {
    137 	return fmt.Sprintf("%s [%s]", usage, key)
    138 }
    139 
    140 func usagePrefixer(key, usage string) string {
    141 	return fmt.Sprintf("[%s] %s", key, usage)
    142 }