From 03fb74b5d738f9ee3a2cf9f99b79643e9e9f1735 Mon Sep 17 00:00:00 2001
From: Chance Zibolski <chance.zibolski@gmail.com>
Date: Fri, 6 Mar 2015 11:21:17 -0800
Subject: [PATCH] Support rewriting env keys

---
 README.md     |  8 +++++++-
 viper.go      | 25 +++++++++++++++++++++++--
 viper_test.go | 14 ++++++++++++++
 3 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 9f236b2..690228c 100644
--- a/README.md
+++ b/README.md
@@ -98,12 +98,13 @@ Aliases permit a single value to be referenced by multiple keys
 ### Working with Environment Variables
 
 Viper has full support for environment variables. This enables 12 factor
-applications out of the box. There are three methods that exist to aid
+applications out of the box. There are four methods that exist to aid
 with working with ENV:
 
  * AutomaticEnv()
  * BindEnv(string...) : error
  * SetEnvPrefix(string)
+ * SetEnvReplacer(string...) *strings.Replacer
 
 _When working with ENV variables it’s important to recognize that Viper
 treats ENV variables as case sensitive._
@@ -131,6 +132,11 @@ any time a viper.Get request is made. It will apply the following rules.
 It will check for a environment variable with a name matching the key
 uppercased and prefixed with the EnvPrefix if set.
 
+SetEnvReplacer allows you to use a `strings.Replacer` object to rewrite Env keys
+to an extent. This is useful if you want to use `-` or something in your Get()
+calls, but want your environmental variables to use `_` delimiters. An example
+of using it can be found in `viper_test.go`
+
 #### Env example
 
 	SetEnvPrefix("spf") // will be uppercased automatically
diff --git a/viper.go b/viper.go
index e808887..2d51b9b 100644
--- a/viper.go
+++ b/viper.go
@@ -77,6 +77,7 @@ type viper struct {
 	envPrefix  string
 
 	automaticEnvApplied bool
+	envKeyReplacer      *strings.Replacer
 
 	config   map[string]interface{}
 	override map[string]interface{}
@@ -153,6 +154,20 @@ func (v *viper) mergeWithEnvPrefix(in string) string {
 	return strings.ToUpper(in)
 }
 
+// TODO: should getEnv logic be moved into find(). Can generalize the use of
+// rewriting keys many things, Ex: Get('someKey') -> some_key
+// (cammel case to snake case for JSON keys perhaps)
+
+// getEnv s a wrapper around os.Getenv which replaces characters in the original
+// key. This allows env vars which have different keys then the config object
+// keys
+func (v *viper) getEnv(key string) string {
+	if v.envKeyReplacer != nil {
+		key = v.envKeyReplacer.Replace(key)
+	}
+	return os.Getenv(key)
+}
+
 // Return the config file used
 func ConfigFileUsed() string            { return v.ConfigFileUsed() }
 func (v *viper) ConfigFileUsed() string { return v.configFile }
@@ -427,7 +442,7 @@ func (v *viper) find(key string) interface{} {
 	if v.automaticEnvApplied {
 		// even if it hasn't been registered, if automaticEnv is used,
 		// check any Get request
-		if val = os.Getenv(v.mergeWithEnvPrefix(key)); val != "" {
+		if val = v.getEnv(v.mergeWithEnvPrefix(key)); val != "" {
 			jww.TRACE.Println(key, "found in environment with val:", val)
 			return val
 		}
@@ -436,7 +451,7 @@ func (v *viper) find(key string) interface{} {
 	envkey, exists := v.env[key]
 	if exists {
 		jww.TRACE.Println(key, "registered as env var", envkey)
-		if val = os.Getenv(envkey); val != "" {
+		if val = v.getEnv(envkey); val != "" {
 			jww.TRACE.Println(envkey, "found in environment with val:", val)
 			return val
 		} else {
@@ -479,6 +494,12 @@ func (v *viper) AutomaticEnv() {
 	v.automaticEnvApplied = true
 }
 
+// SetEnvKeyReplacer sets the strings.Replacer on the viper object
+func SetEnvKeyReplacer(r *strings.Replacer) { v.SetEnvKeyReplacer(r) }
+func (v *viper) SetEnvKeyReplacer(r *strings.Replacer) {
+	v.envKeyReplacer = r
+}
+
 // Aliases provide another accessor for the same key.
 // This enables one to change a name without breaking the application
 func RegisterAlias(alias string, key string) { v.RegisterAlias(alias, key) }
diff --git a/viper_test.go b/viper_test.go
index 2c665ec..a9ba70f 100644
--- a/viper_test.go
+++ b/viper_test.go
@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"os"
 	"sort"
+	"strings"
 	"testing"
 	"time"
 
@@ -221,6 +222,7 @@ func TestEnv(t *testing.T) {
 	AutomaticEnv()
 
 	assert.Equal(t, "crunk", Get("name"))
+
 }
 
 func TestEnvPrefix(t *testing.T) {
@@ -260,6 +262,18 @@ func TestAutoEnvWithPrefix(t *testing.T) {
 	assert.Equal(t, "13", Get("bar"))
 }
 
+func TestSetEnvReplacer(t *testing.T) {
+	Reset()
+
+	AutomaticEnv()
+	os.Setenv("REFRESH_INTERVAL", "30s")
+
+	replacer := strings.NewReplacer("-", "_")
+	SetEnvKeyReplacer(replacer)
+
+	assert.Equal(t, "30s", Get("refresh-interval"))
+}
+
 func TestAllKeys(t *testing.T) {
 	initConfigs()