initial

vendor/local/args/argset.go

@@ -0,0 +1,171 @@
+package args
+
+import (
+	"encoding/json"
+	"errors"
+	"flag"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"strings"
+	"time"
+
+	yaml "gopkg.in/yaml.v2"
+)
+
+type ArgSet struct {
+	parsed      bool
+	args        []*Arg
+	configFiles []string
+}
+
+func NewArgSet(configFiles ...string) *ArgSet {
+	return &ArgSet{
+		args:        []*Arg{},
+		configFiles: configFiles,
+	}
+}
+
+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.setValueFromFiles(); 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 {
+	var err 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)
+		case TIME:
+			arg.Default = fs.String(arg.Flag, arg.Value.(time.Time).Format("2006-01-02"), arg.Help)
+		case FLOAT:
+			arg.Default = fs.Float64(arg.Flag, float64(arg.Value.(float32)), 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)
+		case TIME:
+			arg.Value = *arg.Default.(*string)
+			err = arg.Set(arg.Value)
+		case FLOAT:
+			arg.Value = *arg.Default.(*float64)
+			err = arg.Set(arg.Value)
+		}
+	}
+	return err
+}
+
+func (as *ArgSet) setValueFromFiles() error {
+	for _, file := range as.configFiles {
+		if err := as.setValueFromFile(file); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (as *ArgSet) setValueFromFile(path string) error {
+	b, err := ioutil.ReadFile(path)
+	if err != nil {
+		return err
+	}
+	var config map[string]interface{}
+	if strings.HasSuffix(path, ".json") {
+		err = json.Unmarshal(b, &config)
+	} else if strings.HasSuffix(path, ".yaml") {
+		err = yaml.Unmarshal(b, &config)
+	} else {
+		err = fmt.Errorf("unknown config file suffix: %s", path)
+	}
+	if err != nil {
+		return err
+	}
+	for i := range as.args {
+		key := as.args[i].Flag
+		if v, ok := config[key]; ok {
+			if err := as.args[i].Set(v); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}