From 8dd29d34bf0186945d53ba1ca0cde2324952a6e9 Mon Sep 17 00:00:00 2001
From: brettlangdon <me@brett.is>
Date: Mon, 18 Jan 2016 13:42:04 -0500
Subject: [PATCH] Add support for environment variables

---
 README.md     | 35 +++++++++++++++++++++++++++++++++++
 parse.go      | 17 +++++++++++++++++
 parse_test.go | 50 ++++++++++++++++++++++++++++++++++++++++++++++++++
 usage_test.go |  7 +++++--
 4 files changed, 107 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index f4c8d11..3d1d12f 100644
--- a/README.md
+++ b/README.md
@@ -54,6 +54,41 @@ Input: src.txt
 Output: [x.out y.out z.out]
 ```
 
+### Environment variables
+
+```go
+var args struct {
+	Workers int `arg:"env"`
+}
+arg.MustParse(&args)
+fmt.Println("Workers:", args.Workers)
+```
+
+```
+$ WORKERS=4 ./example
+Workers: 4
+```
+
+```
+$ WORKERS=4 ./example --workers=6
+Workers: 6
+```
+
+You can also override the name of the environment variable:
+
+```go
+var args struct {
+	Workers int `arg:"env:NUM_WORKERS"`
+}
+arg.MustParse(&args)
+fmt.Println("Workers:", args.Workers)
+```
+
+```
+$ NUM_WORKERS=4 ./example
+Workers: 4
+```
+
 ### Usage strings
 ```go
 var args struct {
diff --git a/parse.go b/parse.go
index 219a947..ce3892f 100644
--- a/parse.go
+++ b/parse.go
@@ -18,6 +18,7 @@ type spec struct {
 	required   bool
 	positional bool
 	help       string
+	env        string
 	wasPresent bool
 	isBool     bool
 }
@@ -130,6 +131,13 @@ func NewParser(dests ...interface{}) (*Parser, error) {
 						spec.positional = true
 					case key == "help":
 						spec.help = value
+					case key == "env":
+						// Use override name if provided
+						if value != "" {
+							spec.env = value
+						} else {
+							spec.env = strings.ToUpper(field.Name)
+						}
 					default:
 						return nil, fmt.Errorf("unrecognized tag '%s' on field %s", key, tag)
 					}
@@ -179,6 +187,15 @@ func process(specs []*spec, args []string) error {
 		if spec.short != "" {
 			optionMap[spec.short] = spec
 		}
+		if spec.env != "" {
+			if value, found := os.LookupEnv(spec.env); found {
+				err := setScalar(spec.dest, value)
+				if err != nil {
+					return fmt.Errorf("error processing environment variable %s: %v", spec.env, err)
+				}
+				spec.wasPresent = true
+			}
+		}
 	}
 
 	// process each string from the command line
diff --git a/parse_test.go b/parse_test.go
index 7fca76a..f3e7350 100644
--- a/parse_test.go
+++ b/parse_test.go
@@ -357,3 +357,53 @@ func TestMustParse(t *testing.T) {
 	assert.Equal(t, "bar", args.Foo)
 	assert.NotNil(t, parser)
 }
+
+func TestEnvironmentVariable(t *testing.T) {
+	var args struct {
+		Foo string `arg:"env"`
+	}
+	os.Setenv("FOO", "bar")
+	os.Args = []string{"example"}
+	MustParse(&args)
+	assert.Equal(t, "bar", args.Foo)
+}
+
+func TestEnvironmentVariableOverrideName(t *testing.T) {
+	var args struct {
+		Foo string `arg:"env:BAZ"`
+	}
+	os.Setenv("BAZ", "bar")
+	os.Args = []string{"example"}
+	MustParse(&args)
+	assert.Equal(t, "bar", args.Foo)
+}
+
+func TestEnvironmentVariableOverrideArgument(t *testing.T) {
+	var args struct {
+		Foo string `arg:"env"`
+	}
+	os.Setenv("FOO", "bar")
+	os.Args = []string{"example", "--foo", "baz"}
+	MustParse(&args)
+	assert.Equal(t, "baz", args.Foo)
+}
+
+func TestEnvironmentVariableError(t *testing.T) {
+	var args struct {
+		Foo int `arg:"env"`
+	}
+	os.Setenv("FOO", "bar")
+	os.Args = []string{"example"}
+	err := Parse(&args)
+	assert.Error(t, err)
+}
+
+func TestEnvironmentVariableRequired(t *testing.T) {
+	var args struct {
+		Foo string `arg:"env,required"`
+	}
+	os.Setenv("FOO", "bar")
+	os.Args = []string{"example"}
+	MustParse(&args)
+	assert.Equal(t, "bar", args.Foo)
+}
diff --git a/usage_test.go b/usage_test.go
index 2375e81..255018d 100644
--- a/usage_test.go
+++ b/usage_test.go
@@ -10,9 +10,9 @@ import (
 )
 
 func TestWriteUsage(t *testing.T) {
-	expectedUsage := "usage: example [--name NAME] [--value VALUE] [--verbose] [--dataset DATASET] [--optimize OPTIMIZE] [--ids IDS] INPUT [OUTPUT [OUTPUT ...]]\n"
+	expectedUsage := "usage: example [--name NAME] [--value VALUE] [--verbose] [--dataset DATASET] [--optimize OPTIMIZE] [--ids IDS] [--workers WORKERS] INPUT [OUTPUT [OUTPUT ...]]\n"
 
-	expectedHelp := `usage: example [--name NAME] [--value VALUE] [--verbose] [--dataset DATASET] [--optimize OPTIMIZE] [--ids IDS] INPUT [OUTPUT [OUTPUT ...]]
+	expectedHelp := `usage: example [--name NAME] [--value VALUE] [--verbose] [--dataset DATASET] [--optimize OPTIMIZE] [--ids IDS] [--workers WORKERS] INPUT [OUTPUT [OUTPUT ...]]
 
 positional arguments:
   input
@@ -26,6 +26,8 @@ options:
   --optimize OPTIMIZE, -O OPTIMIZE
                          optimization level
   --ids IDS              Ids
+  --workers WORKERS, -w WORKERS
+                         number of workers to start
   --help, -h             display this help and exit
 `
 	var args struct {
@@ -37,6 +39,7 @@ options:
 		Dataset  string   `arg:"help:dataset to use"`
 		Optimize int      `arg:"-O,help:optimization level"`
 		Ids      []int64  `arg:"help:Ids"`
+		Workers  int      `arg:"-w,env:WORKERS,help:number of workers to start"`
 	}
 	args.Name = "Foo Bar"
 	args.Value = 42