Added methods for working on env slices.

## Changes

* Added methods for working on env slices.
* Added tests
* Load and Overload uses Read (code reuse)
* Avoid using return arguments `raw return` (antipattern)
* Exec does not pollute current process environment variables when invoked.


This allows easy use with with os/exec package.


Signed-off-by: Bartlomiej Plotka <[email protected]>

Author: bwplotka

godotenv.go

@@ -39,16 +39,19 @@ const doubleQuoteSpecialChars = "\\\n\r\"!$`"
 //		godotenv.Load("fileone", "filetwo")
 //
 // 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)
+func Load(filenames ...string) error {
+	envMap, err := Read(filenames...)
+	if err != nil {
+		return err
+	}
 
-	for _, filename := range filenames {
-		err = loadFile(filename, false)
-		if err != nil {
-			return // return early on a spazout
+	for _, e := range MergeEnvSlices(EnvMapToSlice(envMap), os.Environ()) {
+		// We assume env slice always has key and some value (even empty).
+		if err := os.Setenv(strings.Split(e, "=")[0], strings.SplitN(e, "=", 2)[1]); err != nil {
+			return err
 		}
 	}
-	return
+	return nil
 }
 
 // Overload will read your env file(s) and load them into ENV for this process.
@@ -63,15 +66,18 @@ 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) {
-	filenames = filenamesOrDefault(filenames)
+	envMap, err := Read(filenames...)
+	if err != nil {
+		return err
+	}
 
-	for _, filename := range filenames {
-		err = loadFile(filename, true)
-		if err != nil {
-			return // return early on a spazout
+	for _, e := range MergeEnvSlices(os.Environ(), EnvMapToSlice(envMap)) {
+		// We assume env slice always has key and some value (even empty).
+		if err := os.Setenv(strings.Split(e, "=")[0], strings.SplitN(e, "=", 2)[1]); err != nil {
+			return err
 		}
 	}
-	return
+	return nil
 }
 
 // Read all env (with same file loading semantics as Load) but return values as
@@ -82,18 +88,15 @@ func Read(filenames ...string) (envMap map[string]string, err error) {
 
 	for _, filename := range filenames {
 		individualEnvMap, individualErr := readFile(filename)
-
 		if individualErr != nil {
-			err = individualErr
-			return // return early on a spazout
+			return nil, individualErr // Return early on a spazout.
 		}
 
 		for key, value := range individualEnvMap {
 			envMap[key] = value
 		}
 	}
-
-	return
+	return envMap, nil
 }
 
 // Parse reads an env file from io.Reader, returning a map of keys and values.
@@ -107,39 +110,55 @@ func Parse(r io.Reader) (envMap map[string]string, err error) {
 	}
 
 	if err = scanner.Err(); err != nil {
-		return
+		return nil, err
 	}
 
 	for _, fullLine := range lines {
 		if !isIgnoredLine(fullLine) {
 			var key, value string
 			key, value, err = parseLine(fullLine, envMap)
-
 			if err != nil {
-				return
+				return nil, err
 			}
 			envMap[key] = value
 		}
 	}
 	return
 }
 
-//Unmarshal reads an env file from a string, returning a map of keys and values.
+// 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 Parse(strings.NewReader(str))
 }
 
-// Exec loads env vars from the specified filenames (empty map falls back to default)
-// then executes the cmd specified.
+// EnvMapToSlice converts env map that you can get from `Read` or `Parse` into env sorted string
+// slice commonly used by `exec.Command`. This allows to execute new process with relevant
+// variables without changing current process environment.
+// See https://golang.org/pkg/os/exec/#Cmd `Env` field to read more about slice format.
+func EnvMapToSlice(m map[string]string) []string {
+	s := make([]string, 0, len(m))
+	for k, v := range m {
+		s = append(s, fmt.Sprintf("%s=%s", k, v))
+	}
+	sort.Strings(s)
+	return s
+}
+
+// Exec executes given command with args in separate process in environment that consists
+// of env vars from the specified filenames (empty map falls back to default) and current process env vars on top.
 //
 // Simply hooks up os.Stdin/err/out to the command and calls Run()
 //
 // If you want more fine grained control over your command it's recommended
 // that you use `Load()` or `Read()` and the `os/exec` package yourself.
 func Exec(filenames []string, cmd string, cmdArgs []string) error {
-	Load(filenames...)
+	envMap, err := Read(filenames...)
+	if err != nil {
+		return err
+	}
 
 	command := exec.Command(cmd, cmdArgs...)
+	command.Env = MergeEnvSlices(EnvMapToSlice(envMap), os.Environ())
 	command.Stdin = os.Stdin
 	command.Stdout = os.Stdout
 	command.Stderr = os.Stderr
@@ -187,26 +206,41 @@ func filenamesOrDefault(filenames []string) []string {
 	return filenames
 }
 
-func loadFile(filename string, overload bool) error {
-	envMap, err := readFile(filename)
-	if err != nil {
-		return err
-	}
+// MergeEnvSlices merges two slices into single sorted one by applying `over` slice into `base`.
+// The `over` slice will be used if the key overlaps.
+func MergeEnvSlices(base, over []string) (merged []string) {
+	sort.Strings(base)
+	sort.Strings(over)
 
-	currentEnv := map[string]bool{}
-	rawEnv := os.Environ()
-	for _, rawEnvLine := range rawEnv {
-		key := strings.Split(rawEnvLine, "=")[0]
-		currentEnv[key] = true
-	}
+	var b, o int
+	for b < len(base) || o < len(over) {
 
-	for key, value := range envMap {
-		if !currentEnv[key] || overload {
-			os.Setenv(key, value)
+		if b >= len(base) {
+			merged = append(merged, over[o])
+			o++
+			continue
 		}
-	}
 
-	return nil
+		if o >= len(over) {
+			merged = append(merged, base[b])
+			b++
+			continue
+		}
+
+		switch strings.Compare(strings.Split(base[b], "=")[0], strings.Split(over[o], "=")[0]) {
+		case 0:
+			// Same keys. Instead of picking over element, ignore base one. This ensure correct behaviour if base
+			// has duplicate elements.
+			b++
+		case 1:
+			merged = append(merged, over[o])
+			o++
+		case -1:
+			merged = append(merged, base[b])
+			b++
+		}
+	}
+	return merged
 }
 
 func readFile(filename string) (envMap map[string]string, err error) {

godotenv_test.go

@@ -469,6 +469,36 @@ func TestRoundtrip(t *testing.T) {
 		if !reflect.DeepEqual(env, roundtripped) {
 			t.Errorf("Expected '%s' to roundtrip as '%v', got '%v' instead", fixtureFilename, env, roundtripped)
 		}
+	}
+}
 
+func TestEnvMapToSlice(t *testing.T) {
+	expected := []string{
+		"OPTION_A=1",
+		"OPTION_B=2",
+		"OPTION_C=",
+		"OPTION_D=\\n",
+		"OPTION_E=1",
+		"OPTION_F=2",
+		"OPTION_G=",
+		"OPTION_H=\n",
+		"OPTION_I=echo 'asd'",
+		"OPTION_J=postgres://localhost:5432/database?sslmode=disable",
+	}
+
+	converted := EnvMapToSlice(map[string]string{
+		"OPTION_I": "echo 'asd'",
+		"OPTION_A": "1",
+		"OPTION_B": "2",
+		"OPTION_E": "1",
+		"OPTION_C": "",
+		"OPTION_D": "\\n",
+		"OPTION_H": "\n",
+		"OPTION_F": "2",
+		"OPTION_G": "",
+		"OPTION_J": "postgres://localhost:5432/database?sslmode=disable",
+	})
+	if !reflect.DeepEqual(expected, converted) {
+		t.Errorf("Expected '%s' from EnvMapToSlice got '%v' instead", expected, converted)
 	}
 }