From 3c259d20544f284b943b6df306db2f32794ac9b9 Mon Sep 17 00:00:00 2001
From: Aran Wilkinson <aran3001@gmail.com>
Date: Sat, 23 Nov 2024 15:03:43 +0000
Subject: [PATCH] add support for default values with environment variables

---
 README.md                       | 13 ++++++++++--
 config.go                       | 35 ++++++++++++++++++++++++++++++++-
 config_test.go                  |  9 +++++++++
 testdata/test_with_default.yaml |  1 +
 4 files changed, 55 insertions(+), 3 deletions(-)
 create mode 100644 testdata/test_with_default.yaml

diff --git a/README.md b/README.md
index 846ed94..0a57236 100644
--- a/README.md
+++ b/README.md
@@ -1,10 +1,10 @@
 # yamlcfg
 
-yamlcfg is a wrapper around the [gopkg.in/yaml.v3](https://gopkg.in/yaml.v3) library and provides a convient way to configure Golang applications with YAML and environment variables. 
+yamlcfg is a wrapper around the [gopkg.in/yaml.v3](https://gopkg.in/yaml.v3) library and provides a convient way to configure Golang applications with YAML and environment variables.
 
 The library can also automatically call `Validate` functions if present on the given config struct.
 
-## Installation 
+## Installation
 
 To install, run:
 
@@ -66,3 +66,12 @@ func main() {
 }
 
 ```
+
+## Example with default value
+
+
+Example `config.yaml` with environment configuration and default values
+
+```yaml
+log_level: ${LOG_LEVEL:info}
+```
diff --git a/config.go b/config.go
index 43192ec..517c6af 100644
--- a/config.go
+++ b/config.go
@@ -5,6 +5,7 @@ import (
 	"embed"
 	"fmt"
 	"os"
+	"regexp"
 
 	"gopkg.in/yaml.v3"
 )
@@ -49,7 +50,7 @@ func ParseWithConfig[T any](cfg *T, path string) (*T, error) {
 // It will also expand any environment variables in the yaml data
 func UnmarshalConfig[T any](cfg *T, data []byte) error {
 	// expand any $VAR values in the config from environment variables
-	data = []byte(os.ExpandEnv(string(data)))
+	data = []byte(parseEnv(string(data)))
 
 	dec := yaml.NewDecoder(bytes.NewReader(data))
 	dec.KnownFields(true)
@@ -70,3 +71,35 @@ func parse[T any](cfg *T, b []byte) (*T, error) {
 
 	return cfg, nil
 }
+
+// parseEnv replaces $ENV_NAME and ${ENV_NAME:default} placeholders.
+// Default values are only supported with ${ENV_NAME:default}.
+func parseEnv(input string) string {
+	// Regex to match both $ENV_NAME and ${ENV_NAME:default}
+	re := regexp.MustCompile(`\$(\w+)|\$\{(\w+)(?::([^}]*))?\}`)
+
+	return re.ReplaceAllStringFunc(input, func(match string) string {
+		parts := re.FindStringSubmatch(match)
+		if len(parts) == 0 {
+			return match // No match, return as-is
+		}
+
+		// Check if it's $ENV_NAME or ${ENV_NAME:default}
+		if parts[1] != "" {
+			// $ENV_NAME style
+			varName := parts[1]
+			if value, found := os.LookupEnv(varName); found {
+				return value
+			}
+			return "" // If not found, replace with an empty string
+		}
+
+		// ${ENV_NAME:default} style
+		varName := parts[2]
+		defaultValue := parts[3]
+		if value, found := os.LookupEnv(varName); found {
+			return value
+		}
+		return defaultValue // Fallback to default value if not found
+	})
+}
diff --git a/config_test.go b/config_test.go
index 8ddf9e0..2ae5968 100644
--- a/config_test.go
+++ b/config_test.go
@@ -89,6 +89,15 @@ func TestParse(t *testing.T) {
 		qt.Assert(t, qt.Equals(cfg.SomeValue, "this is for testing purposes"))
 	})
 
+	t.Run("successfully parses yaml config with default value", func(t *testing.T) {
+		cfg, err := Parse[TestStruct]("testdata/test_with_default.yaml")
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		qt.Assert(t, qt.Equals(cfg.SomeValue, "this_is_the_default"))
+	})
+
 	t.Run("fails to read unknown file", func(t *testing.T) {
 		cfg, err := Parse[TestStruct]("testdata/this_file_does_not_exist.yaml")
 		qt.Assert(t, qt.IsNotNil(err))
diff --git a/testdata/test_with_default.yaml b/testdata/test_with_default.yaml
new file mode 100644
index 0000000..5473d54
--- /dev/null
+++ b/testdata/test_with_default.yaml
@@ -0,0 +1 @@
+some_value: "${SOME_VALUE:this_is_the_default}"