Merge pull request #6 from ulyssessouza/inherited-vars2

Author: ndeloof

fixtures/inherited-not-found.env

@@ -0,0 +1 @@
+VARIABLE_NOT_FOUND

fixtures/inherited.env

@@ -0,0 +1 @@
+VAR_TO_BE_LOADED_FROM_OS_ENV

godotenv.go

@@ -28,14 +28,26 @@ import (
 
 const doubleQuoteSpecialChars = "\\\n\r\"!$`"
 
+// LookupFn represents a lookup function to resolve variables from
+type LookupFn func(string) (string, bool)
+
+var noLookupFn = func(s string) (string, bool) {
+	return "", false
+}
+
 // Parse reads an env file from io.Reader, returning a map of keys and values.
 func Parse(r io.Reader) (map[string]string, error) {
+	return ParseWithLookup(r, nil)
+}
+
+// ParseWithLookup reads an env file from io.Reader, returning a map of keys and values.
+func ParseWithLookup(r io.Reader, lookupFn LookupFn) (map[string]string, error) {
 	data, err := ioutil.ReadAll(r)
 	if err != nil {
 		return nil, err
 	}
 
-	return UnmarshalBytes(data)
+	return UnmarshalBytesWithLookup(data, lookupFn)
 }
 
 // Load will read your env file(s) and load them into ENV for this process.
@@ -50,15 +62,7 @@ func Parse(r io.Reader) (map[string]string, error) {
 //
 // It's important to note that it WILL NOT OVERRIDE an env variable that already exists - consider the .env file to set dev vars or sensible defaults
 func Load(filenames ...string) (err error) {
-	filenames = filenamesOrDefault(filenames)
-
-	for _, filename := range filenames {
-		err = loadFile(filename, false)
-		if err != nil {
-			return // return early on a spazout
-		}
-	}
-	return
+	return load(false, filenames...)
 }
 
 // Overload will read your env file(s) and load them into ENV for this process.
@@ -73,25 +77,29 @@ func Load(filenames ...string) (err error) {
 //
 // It's important to note this WILL OVERRIDE an env variable that already exists - consider the .env file to forcefilly set all vars.
 func Overload(filenames ...string) (err error) {
+	return load(true, filenames...)
+}
+
+func load(overload bool, filenames ...string) (err error) {
 	filenames = filenamesOrDefault(filenames)
 
 	for _, filename := range filenames {
-		err = loadFile(filename, true)
+		err = loadFile(filename, overload)
 		if err != nil {
 			return // return early on a spazout
 		}
 	}
 	return
 }
 
-// Read all env (with same file loading semantics as Load) but return values as
+// ReadWithLookup gets all env vars from the files and/or lookup function and return values as
 // a map rather than automatically writing values into env
-func Read(filenames ...string) (envMap map[string]string, err error) {
+func ReadWithLookup(lookupFn LookupFn, filenames ...string) (envMap map[string]string, err error) {
 	filenames = filenamesOrDefault(filenames)
 	envMap = make(map[string]string)
 
 	for _, filename := range filenames {
-		individualEnvMap, individualErr := readFile(filename)
+		individualEnvMap, individualErr := readFile(filename, lookupFn)
 
 		if individualErr != nil {
 			err = individualErr
@@ -106,15 +114,26 @@ func Read(filenames ...string) (envMap map[string]string, err error) {
 	return
 }
 
+// Read all env (with same file loading semantics as Load) but return values as
+// a map rather than automatically writing values into env
+func Read(filenames ...string) (envMap map[string]string, err error) {
+	return ReadWithLookup(nil, filenames...)
+}
+
 // Unmarshal reads an env file from a string, returning a map of keys and values.
 func Unmarshal(str string) (envMap map[string]string, err error) {
 	return UnmarshalBytes([]byte(str))
 }
 
 // UnmarshalBytes parses env file from byte slice of chars, returning a map of keys and values.
 func UnmarshalBytes(src []byte) (map[string]string, error) {
+	return UnmarshalBytesWithLookup(src, nil)
+}
+
+// UnmarshalBytesWithLookup parses env file from byte slice of chars, returning a map of keys and values.
+func UnmarshalBytesWithLookup(src []byte, lookupFn LookupFn) (map[string]string, error) {
 	out := make(map[string]string)
-	err := parseBytes(src, out)
+	err := parseBytes(src, out, lookupFn)
 	return out, err
 }
 
@@ -178,7 +197,7 @@ func filenamesOrDefault(filenames []string) []string {
 }
 
 func loadFile(filename string, overload bool) error {
-	envMap, err := readFile(filename)
+	envMap, err := readFile(filename, nil)
 	if err != nil {
 		return err
 	}
@@ -199,14 +218,14 @@ func loadFile(filename string, overload bool) error {
 	return nil
 }
 
-func readFile(filename string) (envMap map[string]string, err error) {
+func readFile(filename string, lookupFn LookupFn) (envMap map[string]string, err error) {
 	file, err := os.Open(filename)
 	if err != nil {
 		return
 	}
 	defer file.Close()
 
-	return Parse(file)
+	return ParseWithLookup(file, lookupFn)
 }
 
 var exportRegex = regexp.MustCompile(`^\s*(?:export\s+)?(.*?)\s*$`)

godotenv_test.go

@@ -487,7 +487,7 @@ func TestRoundtrip(t *testing.T) {
 	fixtures := []string{"equals.env", "exported.env", "plain.env", "quoted.env"}
 	for _, fixture := range fixtures {
 		fixtureFilename := fmt.Sprintf("fixtures/%s", fixture)
-		env, err := readFile(fixtureFilename)
+		env, err := readFile(fixtureFilename, nil)
 		if err != nil {
 			t.Errorf("Expected '%s' to read without error (%v)", fixtureFilename, err)
 		}
@@ -505,3 +505,52 @@ func TestRoundtrip(t *testing.T) {
 
 	}
 }
+
+func TestInheritedEnvVariable(t *testing.T) {
+	const envKey = "VAR_TO_BE_LOADED_FROM_OS_ENV"
+	const envVal = "SOME_RANDOM_VALUE"
+	os.Setenv(envKey, envVal)
+
+	envFileName := "fixtures/inherited.env"
+	expectedValues := map[string]string{
+		envKey: envVal,
+	}
+
+	envMap, err := ReadWithLookup(os.LookupEnv, envFileName)
+	if err != nil {
+		t.Error("Error reading file")
+	}
+	if len(envMap) != len(expectedValues) {
+		t.Error("Didn't get the right size map back")
+	}
+	for key, value := range expectedValues {
+		if envMap[key] != value {
+			t.Error("Read got one of the keys wrong")
+		}
+	}
+}
+
+func TestInheritedEnvVariableNotFound(t *testing.T) {
+	envMap, err := Read("fixtures/inherited-not-found.env")
+	if envMap["VARIABLE_NOT_FOUND"] != "" || err != nil {
+		t.Errorf("Expected 'VARIABLE_NOT_FOUND' to be '' with no errors")
+	}
+}
+
+func TestInheritedEnvVariableNotFoundWithLookup(t *testing.T) {
+	notFoundMap := make(map[string]interface{})
+	envMap, err := ReadWithLookup(func(v string)(string, bool){
+		envVar, ok := os.LookupEnv(v)
+		if !ok {
+			notFoundMap[v] = nil
+		}
+		return envVar, ok
+	}, "fixtures/inherited-not-found.env")
+	if envMap["VARIABLE_NOT_FOUND"] != "" || err != nil {
+		t.Errorf("Expected 'VARIABLE_NOT_FOUND' to be '' with no errors")
+	}
+	_, ok := notFoundMap["VARIABLE_NOT_FOUND"]
+	if !ok {
+		t.Errorf("Expected 'VARIABLE_NOT_FOUND' to be in the set of not found variables")
+	}
+}

parser.go

@@ -16,7 +16,7 @@ const (
 	exportPrefix = "export"
 )
 
-func parseBytes(src []byte, out map[string]string) error {
+func parseBytes(src []byte, out map[string]string, lookupFn LookupFn) error {
 	cutset := src
 	for {
 		cutset = getStatementStart(cutset)
@@ -29,6 +29,22 @@ func parseBytes(src []byte, out map[string]string) error {
 		if err != nil {
 			return err
 		}
+		if strings.Contains(key, " ") {
+			return errors.New("key cannot contain a space")
+		}
+
+		if left == nil  {
+			if lookupFn == nil {
+				lookupFn = noLookupFn
+			}
+
+			value, ok := lookupFn(key)
+			if ok {
+				out[key] = value
+				cutset = left
+				continue
+			}
+		}
 
 		value, left, err := extractVarValue(left, out)
 		if err != nil {
@@ -80,7 +96,7 @@ loop:
 		}
 
 		switch char {
-		case '=', ':':
+		case '=', ':', '\n':
 			// library also supports yaml-style value declaration
 			key = string(src[0:i])
 			offset = i + 1