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 keyLookupType struct {
	key  string
	end  bool
	next map[string]*keyLookupType
}

func (kl *keyLookupType) setKey(parts []string, key string) {
	if len(kl.next) == 0 {
		kl.next = make(map[string]*keyLookupType)
	}
	if len(parts) > 0 {
		firstPart := parts[0]
		nextLookup, ok := kl.next[firstPart]
		if ok {
			nextLookup.setKey(parts[1:], key)
		} else {
			nextLookup := new(keyLookupType)
			nextLookup.setKey(parts[1:], key)
			kl.next[firstPart] = nextLookup
		}
	} else {
		kl.key = key
		kl.end = true
	}
}

func (kl *keyLookupType) print() {
	if kl.end {
		fmt.Println(kl.key)
		fmt.Println("-----------------")
		return
	}
	for key, val := range kl.next {
		fmt.Println(key)
		val.print()
	}
}

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

// 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 = new(mapEnvType)
	for level, _ := range config.mapEnv {
		config.mapEnv[level] = new(keyLookupType)
	}

	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, "_")
			maxParts := len(splitted)
			splitted = append(splitted, make([]string, len(config.mapEnv), len(config.mapEnv))...)
			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 level := 0; level < len(config.mapEnv); level++ {
						for count := 0; count < maxParts-1-level; count++ {
							parts := make([]string, level+2, level+2)
							lastSplit := len(parts) - 1 + count

							for partPos, _ := range parts {
								parts[partPos] = splitted[partPos+count]
							}
							parts[0] = strings.Trim(strings.Join(splitted[:count+1], "_"), "_")
							parts[len(parts)-1] = strings.Trim(strings.Join(splitted[lastSplit:], "_"), "_")
							config.mapEnv[level].setKey(parts, 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[0].next[key]
	if ok {
		for _, entry := range entries.next {
			defEntry := c.env[entry.key]
			defEntry.dtype = dtype
			c.env[entry.key] = defEntry
		}
	}
}

// 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
}

// 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 getFirstRune(str string) rune {
	for _, v := range str {
		return v
	}
	return rune(0)
}