args/argset.go

114 lines
2.1 KiB
Go

package args
import (
"errors"
"flag"
"os"
"time"
)
type ArgSet struct {
parsed bool
args []*Arg
}
func NewArgSet() *ArgSet {
return &ArgSet{
args: []*Arg{},
}
}
func (as *ArgSet) Append(argType Type, key, help string, def interface{}) {
as.args = append(as.args, NewArg(argType, key, help, def))
}
func (as *ArgSet) Get(key string) *Arg {
if !as.parsed {
return nil
}
for i := range as.args {
if as.args[i].Env == key || as.args[i].Flag == key {
return as.args[i]
}
}
return nil
}
func (as *ArgSet) Parse() error {
if as.parsed {
return nil
}
defer func() {
as.parsed = true
}()
if err := as.setValueFromDefaults(); err != nil {
return err
}
if err := as.setValueFromEnv(); err != nil {
return err
}
if err := as.setValueFromFlags(); err != nil {
return err
}
return nil
}
func (as *ArgSet) setValueFromDefaults() error {
for i := range as.args {
if err := as.args[i].Set(as.args[i].Default); err != nil {
return err
}
}
return nil
}
func (as *ArgSet) setValueFromEnv() error {
for i := range as.args {
key := as.args[i].Env
if v, ok := os.LookupEnv(key); ok {
if err := as.args[i].Set(v); err != nil {
return err
}
}
}
return nil
}
func (as *ArgSet) setValueFromFlags() error {
fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
for i := range as.args {
arg := as.args[i]
switch arg.ArgType {
case INT:
arg.Default = fs.Int(arg.Flag, arg.Value.(int), arg.Help)
case STRING:
arg.Default = fs.String(arg.Flag, arg.Value.(string), arg.Help)
case BOOL:
arg.Default = fs.Bool(arg.Flag, arg.Value.(bool), arg.Help)
case DURATION:
arg.Default = fs.Duration(arg.Flag, arg.Value.(time.Duration), arg.Help)
default:
return errors.New("unknown type, cannot set from flag")
}
}
if err := fs.Parse(os.Args[1:]); err != nil {
return err
}
for i := range as.args {
arg := as.args[i]
switch arg.ArgType {
case INT:
arg.Value = *arg.Default.(*int)
case STRING:
arg.Value = *arg.Default.(*string)
case BOOL:
arg.Value = *arg.Default.(*bool)
case DURATION:
arg.Value = *arg.Default.(*time.Duration)
}
}
return nil
}