From 4ee0750c91f1c6eabadd518e9c761a5dd42dbf71 Mon Sep 17 00:00:00 2001
From: Gabriel Aszalos <gabriel.aszalos@gmail.com>
Date: Mon, 14 Sep 2020 11:27:04 +0300
Subject: [PATCH] Allow setting multiple environment variables for the same
 key. (#14)

* Allow setting more than one env. variable into BindEnv.
* Allow setting value transformers which permit the modification of
an environment variable before being assigned to a key.
---
 viper.go      | 47 ++++++++++++++++++++++++++++++++++-------------
 viper_test.go | 34 +++++++++++++++++++++++++++++++++-
 2 files changed, 67 insertions(+), 14 deletions(-)

diff --git a/viper.go b/viper.go
index 21a9a2e..476d5ff 100644
--- a/viper.go
+++ b/viper.go
@@ -194,7 +194,8 @@ type Viper struct {
 	defaults       map[string]interface{}
 	kvstore        map[string]interface{}
 	pflags         map[string]FlagValue
-	env            map[string]string
+	env            map[string][]string
+	envTransform   map[string]func(string) interface{}
 	aliases        map[string]string
 	knownKeys      map[string]interface{}
 	typeByDefValue bool
@@ -217,7 +218,8 @@ func New() *Viper {
 	v.defaults = make(map[string]interface{})
 	v.kvstore = make(map[string]interface{})
 	v.pflags = make(map[string]FlagValue)
-	v.env = make(map[string]string)
+	v.env = make(map[string][]string)
+	v.envTransform = make(map[string]func(string) interface{})
 	v.aliases = make(map[string]string)
 	v.knownKeys = make(map[string]interface{})
 	v.typeByDefValue = false
@@ -368,6 +370,13 @@ func (v *Viper) SetEnvPrefix(in string) {
 	}
 }
 
+// SetEnvKeyTransformer allows defining a transformer function which decides
+// how an environment variables value gets assigned to key.
+func SetEnvKeyTransformer(key string, fn func(string) interface{}) { v.SetEnvKeyTransformer(key, fn) }
+func (v *Viper) SetEnvKeyTransformer(key string, fn func(string) interface{}) {
+	v.envTransform[strings.ToLower(key)] = fn
+}
+
 func (v *Viper) mergeWithEnvPrefix(in string) string {
 	if v.envPrefix != "" {
 		return strings.ToUpper(v.envPrefix + "_" + in)
@@ -1023,21 +1032,20 @@ func (v *Viper) BindFlagValue(key string, flag FlagValue) error {
 // EnvPrefix will be used when set when env name is not provided.
 func BindEnv(input ...string) error { return v.BindEnv(input...) }
 func (v *Viper) BindEnv(input ...string) error {
-	var key, envkey string
 	if len(input) == 0 {
 		return fmt.Errorf("BindEnv missing key to bind to")
 	}
 
-	key = strings.ToLower(input[0])
+	key := strings.ToLower(input[0])
+	var envkeys []string
 
 	if len(input) == 1 {
-		envkey = v.mergeWithEnvPrefix(key)
+		envkeys = []string{v.mergeWithEnvPrefix(key)}
 	} else {
-		envkey = input[1]
+		envkeys = input[1:]
 	}
 
-	v.env[key] = envkey
-
+	v.env[key] = append(v.env[key], envkeys...)
 	v.SetKnown(key)
 
 	return nil
@@ -1108,10 +1116,15 @@ func (v *Viper) find(lcaseKey string) interface{} {
 			return nil
 		}
 	}
-	envkey, exists := v.env[lcaseKey]
+	envkeys, exists := v.env[lcaseKey]
 	if exists {
-		if val, ok := v.getEnv(envkey); ok {
-			return val
+		for _, key := range envkeys {
+			if val, ok := v.getEnv(key); ok {
+				if fn, ok := v.envTransform[lcaseKey]; ok {
+					return fn(val)
+				}
+				return val
+			}
 		}
 	}
 	if nested && v.isPathShadowedInFlatMap(path, v.env) != "" {
@@ -1629,6 +1642,14 @@ func castToMapStringInterface(
 	return tgt
 }
 
+func castMapStringSliceToMapInterface(src map[string][]string) map[string]interface{} {
+	tgt := map[string]interface{}{}
+	for k, v := range src {
+		tgt[k] = v
+	}
+	return tgt
+}
+
 func castMapStringToMapInterface(src map[string]string) map[string]interface{} {
 	tgt := map[string]interface{}{}
 	for k, v := range src {
@@ -1795,7 +1816,7 @@ func (v *Viper) AllKeys() []string {
 	m = v.flattenAndMergeMap(m, castMapStringToMapInterface(v.aliases), "")
 	m = v.flattenAndMergeMap(m, v.override, "")
 	m = v.mergeFlatMap(m, castMapFlagToMapInterface(v.pflags))
-	m = v.mergeFlatMap(m, castMapStringToMapInterface(v.env))
+	m = v.mergeFlatMap(m, castMapStringSliceToMapInterface(v.env))
 	m = v.flattenAndMergeMap(m, v.config, "")
 	m = v.flattenAndMergeMap(m, v.kvstore, "")
 	m = v.flattenAndMergeMap(m, v.defaults, "")
@@ -1850,7 +1871,7 @@ func (v *Viper) flattenAndMergeMap(shadow map[string]bool, m map[string]interfac
 func (v *Viper) mergeFlatMap(shadow map[string]bool, m map[string]interface{}) map[string]bool {
 	// scan keys
 outer:
-	for k, _ := range m {
+	for k := range m {
 		path := strings.Split(k, v.keyDelim)
 		// scan intermediate paths
 		var parentKey string
diff --git a/viper_test.go b/viper_test.go
index 912c6de..21efc4b 100644
--- a/viper_test.go
+++ b/viper_test.go
@@ -383,10 +383,11 @@ func TestEnv(t *testing.T) {
 	initJSON(v)
 
 	v.BindEnv("id")
-	v.BindEnv("f", "FOOD")
+	v.BindEnv("f", "FOOD", "DEPRECATED_FOOD")
 
 	os.Setenv("ID", "13")
 	os.Setenv("FOOD", "apple")
+	os.Setenv("DEPRECATED_FOOD", "banana")
 	os.Setenv("NAME", "crunk")
 
 	assert.Equal(t, "13", v.Get("id"))
@@ -396,7 +397,38 @@ func TestEnv(t *testing.T) {
 	v.AutomaticEnv()
 
 	assert.Equal(t, "crunk", v.Get("name"))
+}
 
+func TestEnvTransformer(t *testing.T) {
+	v := New()
+	initJSON(v)
+
+	v.BindEnv("id", "MY_ID")
+	v.SetEnvKeyTransformer("id", func(in string) interface{} { return "transformed" })
+	os.Setenv("MY_ID", "14")
+
+	assert.Equal(t, "transformed", v.Get("id"))
+}
+
+func TestMultipleEnv(t *testing.T) {
+	v := New()
+	initJSON(v)
+
+	v.BindEnv("id")
+	v.BindEnv("f", "FOOD", "DEPRECATED_FOOD")
+
+	os.Setenv("ID", "13")
+	os.Unsetenv("FOOD")
+	os.Setenv("DEPRECATED_FOOD", "banana")
+	os.Setenv("NAME", "crunk")
+
+	assert.Equal(t, "13", v.Get("id"))
+	assert.Equal(t, "banana", v.Get("f"))
+	assert.Equal(t, "Cake", v.Get("name"))
+
+	v.AutomaticEnv()
+
+	assert.Equal(t, "crunk", v.Get("name"))
 }
 
 func TestEmptyEnv(t *testing.T) {