Merge pull request #9 from tbocek/master

Author: ndeloof

godotenv.go

@@ -231,6 +231,9 @@ func readFile(filename string, lookupFn LookupFn) (envMap map[string]string, err
 var exportRegex = regexp.MustCompile(`^\s*(?:export\s+)?(.*?)\s*$`)
 
 func parseLine(line string, envMap map[string]string) (key string, value string, err error) {
+	return parseLineWithLookup(line, envMap, nil)
+}
+func parseLineWithLookup(line string, envMap map[string]string, lookupFn LookupFn) (key string, value string, err error) {
 	if len(line) == 0 {
 		err = errors.New("zero length string")
 		return
@@ -274,7 +277,7 @@ func parseLine(line string, envMap map[string]string) (key string, value string,
 	key = exportRegex.ReplaceAllString(splitString[0], "$1")
 
 	// Parse the value
-	value = parseValue(splitString[1], envMap)
+	value = parseValue(splitString[1], envMap, lookupFn)
 	return
 }
 
@@ -285,7 +288,7 @@ var (
 	unescapeCharsRegex = regexp.MustCompile(`\\([^$])`)
 )
 
-func parseValue(value string, envMap map[string]string) string {
+func parseValue(value string, envMap map[string]string, lookupFn LookupFn) string {
 
 	// trim
 	value = strings.Trim(value, " ")
@@ -319,7 +322,7 @@ func parseValue(value string, envMap map[string]string) string {
 		}
 
 		if singleQuotes == nil {
-			value = expandVariables(value, envMap)
+			value = expandVariables(value, envMap, lookupFn)
 		}
 	}
 
@@ -328,7 +331,7 @@ func parseValue(value string, envMap map[string]string) string {
 
 var expandVarRegex = regexp.MustCompile(`(\\)?(\$)(\()?\{?([A-Z0-9_]+)?\}?`)
 
-func expandVariables(v string, m map[string]string) string {
+func expandVariables(v string, envMap map[string]string, lookupFn LookupFn) string {
 	return expandVarRegex.ReplaceAllStringFunc(v, func(s string) string {
 		submatch := expandVarRegex.FindStringSubmatch(s)
 
@@ -338,7 +341,20 @@ func expandVariables(v string, m map[string]string) string {
 		if submatch[1] == "\\" || submatch[2] == "(" {
 			return submatch[0][1:]
 		} else if submatch[4] != "" {
-			return m[submatch[4]]
+			//first check if we have defined this already earlier
+			if envMap[submatch[4]] != "" {
+				return envMap[submatch[4]]
+			}
+			if lookupFn == nil {
+				return ""
+			}
+			//if we have not defined it, check the lookup function provided
+			//by the user
+			s2, ok := lookupFn(submatch[4])
+			if ok {
+				return s2
+			}
+			return ""
 		}
 		return s
 	})

godotenv_test.go

@@ -514,8 +514,8 @@ func TestInheritedEnvVariablSameSize(t *testing.T) {
 	envFileName := "fixtures/inherited-multi-var.env"
 	expectedValues := map[string]string{
 		envKey: envVal,
-		"foo": "bar",
-		"bar": "baz",
+		"foo":  "bar",
+		"bar":  "baz",
 	}
 
 	envMap, err := ReadWithLookup(os.LookupEnv, envFileName)
@@ -565,7 +565,7 @@ func TestInheritedEnvVariableNotFound(t *testing.T) {
 
 func TestInheritedEnvVariableNotFoundWithLookup(t *testing.T) {
 	notFoundMap := make(map[string]interface{})
-	envMap, err := ReadWithLookup(func(v string)(string, bool){
+	envMap, err := ReadWithLookup(func(v string) (string, bool) {
 		envVar, ok := os.LookupEnv(v)
 		if !ok {
 			notFoundMap[v] = nil
@@ -579,4 +579,18 @@ func TestInheritedEnvVariableNotFoundWithLookup(t *testing.T) {
 	if !ok {
 		t.Errorf("Expected 'VARIABLE_NOT_FOUND' to be in the set of not found variables")
 	}
-}
+}
+
+func TestExpendingEnvironmentWithLookup(t *testing.T) {
+	rawEnvLine := "TEST=$ME"
+	expectedValue := "YES"
+	key, value, _ := parseLineWithLookup(rawEnvLine, noopPresets, func(s string) (string, bool) {
+		if s == "ME" {
+			return expectedValue, true
+		}
+		return "NO", false
+	})
+	if value != "YES" {
+		t.Errorf("Expected '%v' to parse as '%v' => '%v', got '%v' => '%v' instead", rawEnvLine, key, expectedValue, key, value)
+	}
+}

parser.go

@@ -46,7 +46,7 @@ func parseBytes(src []byte, out map[string]string, lookupFn LookupFn) error {
 			}
 		}
 
-		value, left, err := extractVarValue(left, out)
+		value, left, err := extractVarValue(left, out, lookupFn)
 		if err != nil {
 			return err
 		}
@@ -126,16 +126,16 @@ loop:
 }
 
 // extractVarValue extracts variable value and returns rest of slice
-func extractVarValue(src []byte, vars map[string]string) (value string, rest []byte, err error) {
+func extractVarValue(src []byte, envMap map[string]string, lookupFn LookupFn) (value string, rest []byte, err error) {
 	quote, hasPrefix := hasQuotePrefix(src)
 	if !hasPrefix {
 		// unquoted value - read until whitespace
 		end := bytes.IndexFunc(src, unicode.IsSpace)
 		if end == -1 {
-			return expandVariables(string(src), vars), nil, nil
+			return expandVariables(string(src), envMap, lookupFn), nil, nil
 		}
 
-		return expandVariables(string(src[0:end]), vars), src[end:], nil
+		return expandVariables(string(src[0:end]), envMap, lookupFn), src[end:], nil
 	}
 
 	// lookup quoted string terminator
@@ -155,7 +155,7 @@ func extractVarValue(src []byte, vars map[string]string) (value string, rest []b
 		if quote == prefixDoubleQuote {
 			// unescape newlines for double quote (this is compat feature)
 			// and expand environment variables
-			value = expandVariables(expandEscapes(value), vars)
+			value = expandVariables(expandEscapes(value), envMap, lookupFn)
 		}
 
 		return value, src[i+1:], nil