envconf/envconf.go

package envconf

import (
	"errors"
	"fmt"
	"os"
	"regexp"
	"strings"
	"time"
	"unicode"
)

type cEntry struct {
	value  string
	parsed cValue
	dtype  DataType
	unset  bool
	empty  bool
	defval string // The default value
	hasdef bool   // Default value is defined
}

type Config struct {
	parsed bool
	env    map[string]cEntry
	mapEnv map[string]map[string]string
}

// NewConfig returns an envconf.Config that is used to read configuration from environment variables.
// The environment variables are stored in envconf.Config, so changes to the environment after NewConfig has been called
// will not be taken into account.
func NewConfig() *Config {
	config := new(Config)
	config.parsed = false
	config.env = make(map[string]cEntry)
	config.mapEnv = make(map[string]map[string]string)

	for _, v := range os.Environ() {
		splitted := strings.SplitN(v, "=", 2)
		if len(splitted) == 2 {
			key := cleanKey(splitted[0])
			val := splitted[1]
			splitted = strings.Split(key, "_")
			if unicode.IsLetter(getFirstRune(key)) {
				var entry cEntry
				entry.value = val
				entry.dtype = TypeNone
				entry.unset = false
				entry.empty = false
				config.env[key] = entry
				if len(splitted) > 1 {
					for count, _ := range splitted {
						left := strings.Join(splitted[:count], "_")
						right := strings.Join(splitted[count:], "_")
						if len(config.mapEnv[left]) == 0 {
							config.mapEnv[left] = make(map[string]string)
						}
						config.mapEnv[left][right] = key
					}
				}
			}
		}
	}
	return config
}

// Define the type of an environment variable.
// Variables without a defined type will be ignored by Parse.
func (c *Config) Define(key string, dtype DataType) {
	key = cleanKey(key)
	entry, ok := c.env[key]
	if ok {
		entry.dtype = dtype
		c.env[key] = entry
	} else {
		var entry cEntry
		entry.dtype = dtype
		entry.unset = true
		entry.empty = true
		c.env[key] = entry
	}
}

// Define the type of an environment variable.
// Variables without a defined type will be ignored by Parse.
func (c *Config) DefineMap(key string, dtype DataType) {
	key = cleanKey(key)
	entries, ok := c.mapEnv[key]
	if ok {
		for _, key = range entries {
			entry := c.env[key]
			entry.dtype = dtype
			c.env[key] = entry
		}
	}
}

// Define the type and default value of an environment variable.
// Variables without a defined type will be ignored by Parse.
func (c *Config) DefineDefault(key string, val string, dtype DataType) {
	key = cleanKey(key)
	entry, ok := c.env[key]
	if ok {
		if entry.unset {
			entry.value = val
		}
		entry.dtype = dtype
		entry.empty = false
		entry.defval = val
		entry.hasdef = true
		c.env[key] = entry
	} else {
		var entry cEntry
		entry.dtype = dtype
		entry.unset = true
		entry.empty = false
		entry.value = val
		entry.defval = val
		entry.hasdef = true
		c.env[key] = entry
	}
}

// Parse parses the environment variables previously defined by Define and DefineDefault.
// Parse should only be called once for a given envconf.Config.
func (c *Config) Parse() {
	if c.parsed {
		return
	}
	c.parsed = true
	failed := false
	for k, v := range c.env {
		if v.empty {
			if v.unset {
				failed = true
			}
		} else {
			v.parsed = v.dtype.parse(k, v.value)
			c.env[k] = v
		}
	}

	if failed {
		for k, v := range c.env {
			if (v.parsed.err == nil) && v.unset {
				if v.empty {
					v.parsed.err = errors.New(fmt.Sprintf(`Environment variable "%s" not found. It should have been of type %s.`, k, v.dtype))
					c.env[k] = v
				} else {
					v.parsed.err = errors.New(fmt.Sprintf(`Environment variable "%s" not found. Default value "%s" of type %s used.`, k, v.value, v.dtype))
					c.env[k] = v
				}
			}
		}
	}
}

func (c *Config) help() {
	max := make([]int, 2, 2)
	preq := false
	pdef := false
	for k, v := range c.env {
		if v.dtype != TypeNone {
			if len(k) > max[0] {
				max[0] = len(k)
			}
			if len(v.dtype.String()) > max[1] {
				max[1] = len(v.dtype.String())
			}
			if v.hasdef {
				pdef = true
			} else {
				preq = true
			}
		}
	}
	if pdef {
		fmt.Println()
		for k, v := range c.env {
			if v.dtype != TypeNone {
				if v.hasdef {
					format := fmt.Sprintf("Variable   %%-%ds|   Type %%-%ds|   Default %%s\n", max[0]+5, max[1]+3)
					fmt.Printf(format, fmt.Sprintf(`"%s"`, k), v.dtype, fmt.Sprintf(`"%s"`, v.defval))
				}
			}
		}
	}
	if preq {
		fmt.Println()
		for k, v := range c.env {
			if v.dtype != TypeNone {
				if !v.hasdef {
					format := fmt.Sprintf("Variable   %%-%ds|   Type %%-%ds|   Required\n", max[0]+5, max[1]+3)
					fmt.Printf(format, fmt.Sprintf(`"%s"`, k), v.dtype)
				}
			}
		}
	}
	fmt.Println()
}

func (c *Config) StatusHelp() bool {
	flags := make(map[string]bool)
	flags["--help"] = true
	flags["-help"] = true
	flags["-h"] = true
	if len(os.Args) > 1 {
		key := strings.ToLower(strings.TrimSpace(os.Args[1]))
		_, ok := flags[key]
		if ok {
			c.help()
			return false
		}
	}
	return c.Status()
}

// Status prints out failures that occured while parsing the environment to os.Stderr.
// Variables that have been defined without a default value and are
// missing from the environment will be considered a failure.
// If parsing of any of the variables has failed Status will return false.
func (c *Config) Status() (ok bool) {
	ok = c.parsed
	if ok {
		for _, v := range c.env {
			err := v.parsed.err
			if err != nil {
				ok = false
				if !v.empty {
					fmt.Fprintln(os.Stderr, err)
				}
			}
		}

		if !ok {
			fmt.Fprintln(os.Stderr, "")
			for _, v := range c.env {
				err := v.parsed.err
				if (err != nil) && v.empty {
					fmt.Fprintln(os.Stderr, err)
				}
			}
		}
	}
	return
}

func cleanKey(key string) string {
	expr := regexp.MustCompile("__+")
	fn := func(r rune) rune {
		if (r >= '0' && r <= '9') || (r >= 'A' && r <= 'Z') || r == '_' {
			return r
		}
		return -1
	}

	key = strings.Trim(strings.Map(fn, strings.ToUpper(key)), "_")
	return expr.ReplaceAllString(key, "_")
}

func keySplit(key string) (left string, right string, ok bool) {
	key = cleanKey(key)
	pos := strings.LastIndex(key, "_")
	ok = false
	if (pos+1 < len(key)) && (pos > 1) {
		return cleanKey(key[:pos]), cleanKey(key[pos:]), true
	}
	return
}

func (c *Config) getRaw(key string, dtype DataType) (val cValue) {
	val.dtype = TypeNone
	val.binval = make([]byte, 0, 0)
	if c.parsed {
		key = cleanKey(key)
		entry, ok := c.env[key]
		if ok && (entry.dtype.baseType() == dtype.baseType()) {
			return entry.parsed
		}
	}
	return
}

func (c *Config) getRawMap(key string, dtype DataType) (empty map[string]cValue) {
	empty = make(map[string]cValue)
	retval := make(map[string]cValue)
	if c.parsed {
		key = cleanKey(key)
		entries, ok := c.mapEnv[key]

		if ok {
			for k, v := range entries {
				entry := c.env[v]
				if (entry.dtype.baseType() == dtype.baseType()) && (entry.parsed.err == nil) {
					retval[k] = entry.parsed
				} else {
					return
				}
			}
			return retval
		}
	}
	return
}

// Returns the value of an environment variable.
// If the variable is not defined as envconf.TypeInt the function will return 0.
func (c *Config) GetInt(key string) int64 {
	val := c.getRaw(key, TypeInt)
	return val.intval
}

// Returns the value of an environment variable.
// If the variable is not defined as envconf.TypeHex the function will return []byte{}.
func (c *Config) GetHex(key string) []byte {
	val := c.getRaw(key, TypeHex)
	return val.binval
}

// Returns the value of an environment variable.
// If the variable is not defined as envconf.TypeMetric the function will return 0.
func (c *Config) GetMetric(key string) int64 {
	val := c.getRaw(key, TypeMetric)
	return val.intval
}

// Returns the value of an environment variable.
// If the variable is not defined as envconf.TypeDirectory the
// function will return the empty string.
func (c *Config) GetDirectory(key string) string {
	val := c.getRaw(key, TypeDirectory)
	return val.strval
}

// Returns the value of an environment variable.
// If the variable is not defined as envconf.TypeString the
// function will return the empty string.
func (c *Config) GetString(key string) string {
	val := c.getRaw(key, TypeString)
	return val.strval
}

// Returns the value of an environment variable.
// If the variable is not defined as envconf.TypeDuration the
// function will return time.Duration(0).
func (c *Config) GetDuration(key string) time.Duration {
	val := c.getRaw(key, TypeDuration)
	return val.durval
}

// Returns the value of an environment variable.
// If the variable is not defined as envconf.TypeBool the
// function will return false.
func (c *Config) GetBool(key string) bool {
	val := c.getRaw(key, TypeBool)
	return val.boolval
}

func (c *Config) GetMapInt(key string) (retval map[string]int64) {
	retval = make(map[string]int64)
	for k, v := range c.getRawMap(key, TypeInt) {
		retval[k] = v.intval
	}
	return
}

func (c *Config) GetMapDuration(key string) (retval map[string]time.Duration) {
	retval = make(map[string]time.Duration)
	for k, v := range c.getRawMap(key, TypeDuration) {
		retval[k] = v.durval
	}
	return
}

func (c *Config) GetMapString(key string) (retval map[string]string) {
	retval = make(map[string]string)
	for k, v := range c.getRawMap(key, TypeString) {
		retval[k] = v.strval
	}
	return
}

func (c *Config) GetMapBool(key string) (retval map[string]bool) {
	retval = make(map[string]bool)
	for k, v := range c.getRawMap(key, TypeBool) {
		retval[k] = v.boolval
	}
	return
}

func (c *Config) GetMapHex(key string) (retval map[string][]byte) {
	retval = make(map[string][]byte)
	for k, v := range c.getRawMap(key, TypeHex) {
		retval[k] = v.binval
	}
	return
}

func getFirstRune(str string) rune {
	for _, v := range str {
		return v
	}
	return rune(0)
}