Start envflag

6 files changed, 262 insertions, 0 deletions

diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..0857be3
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2022 Felix Hanley
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
+of the Software, and to permit persons to whom the Software is furnished to do
+so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..978b8af
--- /dev/null
+++ b/README.md
@@ -0,0 +1,3 @@
+# Environment variables and flags
+
+Set flags using Go's `flag` package, and complement them with environment variables.
diff --git a/example.go b/example.go
new file mode 100644
index 0000000..5bffebd
--- /dev/null
+++ b/example.go
@@ -0,0 +1,17 @@
+package envflag
+
+import (
+	"flag"
+	"os"
+)
+
+func ExampleFlag() {
+	flag.Bool("verbose", false, "Be verbose")
+	Parse()
+}
+
+func ExampleFlagSet() {
+	fs := flag.NewFlagSet("example", flag.ExitOnError)
+	fs.Bool("verbose", false, "Be verbose")
+	ParseFlagSet(fs, os.Args[1:])
+}
diff --git a/flag.go b/flag.go
new file mode 100644
index 0000000..9e12188
--- /dev/null
+++ b/flag.go
@@ -0,0 +1,119 @@
+package envflag
+
+import (
+	"flag"
+	"fmt"
+	"os"
+	"regexp"
+	"strings"
+)
+
+// Option type for configing the parsing.
+type Option option
+type option func(*config)
+
+// FlagConverterFunc functions are used to update flag name for envvar lookups.
+type FlagConverterFunc func(flag string) string
+
+// UsageUpdaterFunc functions are used to update flag usage strings.
+type UsageUpdaterFunc func(key, usage string) string
+
+type config struct {
+	flagConverter FlagConverterFunc
+	usageUpdater  UsageUpdaterFunc
+}
+
+// UsageUpdater enables configurable flag usage updates.
+func UsageUpdater(f UsageUpdaterFunc) Option {
+	return func(cfg *config) {
+		cfg.usageUpdater = f
+	}
+}
+
+// UsagePrefixer prefixes the flag usage with [envvar].
+func UsagePrefixer() Option {
+	return func(cfg *config) {
+		cfg.usageUpdater = usagePrefixer
+	}
+}
+
+// UsageSuffixer suffixes the flag usage with [envvar].
+func UsageSuffixer() Option {
+	return func(cfg *config) {
+		cfg.usageUpdater = usageSuffixer
+	}
+}
+
+// FlagConverter converts the flag name to an environment variable key.
+func FlagConverter(f FlagConverterFunc) Option {
+	return func(cfg *config) {
+		cfg.flagConverter = f
+	}
+}
+
+// ParseWithEnv parses the command-line flags from os.Args[1:]. Must be called
+// after all flags are defined and before flags are accessed by the program.
+func Parse(opts ...Option) error {
+	return parseFlagSetWithEnv(flag.CommandLine, os.Args[1:], opts...)
+}
+
+// ParseFlagSetWithEnv parses flag definitions from the argument list, which
+// should not include the command name. Must be called after all flags in the
+// FlagSet are defined and before flags are accessed by the program. The return
+// value will be ErrHelp if -help or -h were set but not defined.
+func ParseFlagSet(fs *flag.FlagSet, arguments []string, opts ...Option) error {
+	return parseFlagSetWithEnv(fs, arguments, opts...)
+}
+
+func parseFlagSetWithEnv(fs *flag.FlagSet, arguments []string, opts ...Option) error {
+	if fs.Parsed() {
+		return fmt.Errorf("flag has already been parsed")
+	}
+	cfg := &config{
+		flagConverter: flagToEnv(),
+	}
+	for _, o := range opts {
+		o(cfg)
+	}
+	var nerr error
+	fs.VisitAll(func(f *flag.Flag) {
+		envKey := cfg.flagConverter(f.Name)
+		if cfg.usageUpdater != nil {
+			f.Usage = cfg.usageUpdater(envKey, f.Usage)
+		}
+		if envVal := os.Getenv(envKey); envVal != "" {
+			err := f.Value.Set(envVal)
+			if err != nil && nerr == nil {
+				nerr = err
+			}
+		}
+	})
+	if nerr != nil {
+		return nerr
+	}
+	return fs.Parse(arguments)
+}
+
+// convert this-format to THIS_FORMAT
+func flagToEnv() func(string) string {
+	re := regexp.MustCompile("[^a-zA-Z0-9_]+")
+	return func(f string) string {
+		return strings.ToUpper(re.ReplaceAllLiteralString(f, "_"))
+	}
+}
+
+func usageSuffixer(key, usage string) string {
+	envSuffix := fmt.Sprintf("[%s]", key)
+	if strings.HasSuffix(usage, envSuffix) {
+		return usage
+	}
+	return fmt.Sprintf("%s %s", usage, envSuffix)
+}
+
+func usagePrefixer(key, usage string) string {
+	envPrefix := fmt.Sprintf("[%s]", key)
+	if strings.HasPrefix(usage, envPrefix) {
+		return usage
+	}
+	return fmt.Sprintf("%s %s", envPrefix, usage)
+}
diff --git a/flag_test.go b/flag_test.go
new file mode 100644
index 0000000..bed8d8d
--- /dev/null
+++ b/flag_test.go
@@ -0,0 +1,99 @@
+package envflag
+
+import (
+	"flag"
+	"os"
+	"testing"
+)
+
+func TestParseFlagSet(t *testing.T) {
+	tests := map[string]struct {
+		// Flag name
+		name string
+		// Default value
+		def string
+		// Flag command values
+		values []string
+		opts   []Option
+		// The environment
+		env map[string]string
+
+		expected string
+		usage    string
+	}{
+		"no flags no env default": {
+			name:     "testing-123",
+			def:      "defaultvalue",
+			expected: "defaultvalue",
+		},
+		"flags no env no default": {
+			name:     "testing-123",
+			values:   []string{"-testing-123", "passed"},
+			expected: "passed",
+		},
+		"no flags env": {
+			name:     "testing-123",
+			env:      map[string]string{"TESTING_123": "fromenv"},
+			expected: "fromenv",
+		},
+		"flags with no env": {
+			name:     "testing-123",
+			values:   []string{"-testing-123", "passed"},
+			expected: "passed",
+		},
+		"flags with env": {
+			// Flags take precedence
+			name:     "testing-123",
+			values:   []string{"-testing-123", "passed"},
+			env:      map[string]string{"TESTING_123": "fromenv"},
+			expected: "passed",
+		},
+		"no flags with env": {
+			// Flags take precedence
+			name:     "testing-123",
+			env:      map[string]string{"TESTING_123": "fromenv"},
+			expected: "fromenv",
+		},
+		"flags with env and suffix": {
+			name:     "testing-123",
+			values:   []string{"-testing-123", "passed"},
+			opts:     []Option{UsageSuffixer()},
+			env:      map[string]string{"TESTING_123": "fromenv"},
+			expected: "passed",
+		},
+		"no flags with env and suffix": {
+			name:     "testing-123",
+			def:      "defaultvalue",
+			opts:     []Option{UsageSuffixer()},
+			env:      map[string]string{"TESTING_123": "fromenv"},
+			expected: "fromenv",
+		},
+		"a variety of characters": {
+			// Flags take precedence
+			name:     "tes_t.ing----123",
+			env:      map[string]string{"TES_T_ING_123": "fromenv"},
+			expected: "fromenv",
+		},
+	}
+
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			fs := flag.NewFlagSet("test", flag.ExitOnError)
+			for k, v := range tt.env {
+				os.Setenv(k, v)
+			}
+			defer func() {
+				for k := range tt.env {
+					os.Unsetenv(k)
+				}
+			}()
+			actual := fs.String(tt.name, tt.def, "usage")
+			ParseFlagSet(fs, tt.values, tt.opts...)
+
+			if *actual != tt.expected {
+				t.Errorf("got %q, want %q", *actual, tt.expected)
+			}
+		})
+	}
+
+}
diff --git a/go.mod b/go.mod
new file mode 100644
index 0000000..c91313e
--- /dev/null
+++ b/go.mod
@@ -0,0 +1,3 @@
+module userspace.com.au/envflag
+
+go 1.19