From 2770a34ce6d4b41741a2fc65ca666868f1c55fa2 Mon Sep 17 00:00:00 2001
From: Dan <artifactori62@gmail.com>
Date: Fri, 22 Oct 2021 11:28:58 +0700
Subject: [PATCH] Add envprefix tag

---
 README.md        |  1 +
 build.go         |  3 +++
 help_test.go     | 14 ++++++++++++++
 resolver_test.go | 33 +++++++++++++++++++++++++++++++++
 tag.go           |  2 ++
 5 files changed, 53 insertions(+)

diff --git a/README.md b/README.md
index 57ec0a8..a065193 100644
--- a/README.md
+++ b/README.md
@@ -450,6 +450,7 @@ Tag                    | Description
 `group:"X"`            | Logical group for a flag or command.
 `xor:"X,Y,..."`        | Exclusive OR groups for flags. Only one flag in the group can be used which is restricted within the same command. When combined with `required`, at least one of the `xor` group will be required.
 `prefix:"X"`           | Prefix for all sub-flags.
+`envprefix:"X"`        | Envar prefix for all sub-flags.
 `set:"K=V"`            | Set a variable for expansion by child elements. Multiples can occur.
 `embed:""`             | If present, this field's children will be embedded in the parent. Useful for composition.
 `passthrough:""`       | If present, this positional argument stops flag parsing when encountered, as if `--` was processed before. Useful for external command wrappers, like `exec`.
diff --git a/build.go b/build.go
index 95ca7e1..aa5a143 100644
--- a/build.go
+++ b/build.go
@@ -97,6 +97,7 @@ func flattenedFields(v reflect.Value) (out []flattenedField, err error) {
 			}
 			// Accumulate prefixes.
 			subf.tag.Prefix = tag.Prefix + subf.tag.Prefix
+			subf.tag.EnvPrefix = tag.EnvPrefix + subf.tag.EnvPrefix
 			// Combine parent vars.
 			subf.tag.Vars = tag.Vars.CloneWith(subf.tag.Vars)
 		}
@@ -138,6 +139,8 @@ MAIN:
 			name = tag.Prefix + name
 		}
 
+		tag.Env = tag.EnvPrefix + tag.Env
+
 		// Nested structs are either commands or args, unless they implement the Mapper interface.
 		if field.value.Kind() == reflect.Struct && (tag.Cmd || tag.Arg) && k.registry.ForValue(fv) == nil {
 			typ := CommandNode
diff --git a/help_test.go b/help_test.go
index 728681e..875ca1b 100644
--- a/help_test.go
+++ b/help_test.go
@@ -461,6 +461,20 @@ func TestEnvarAutoHelp(t *testing.T) {
 	require.Contains(t, w.String(), "A flag ($FLAG).")
 }
 
+func TestEnvarAutoHelpWithEnvPrefix(t *testing.T) {
+	type Anonymous struct {
+		Flag string `env:"FLAG" help:"A flag."`
+	}
+	var cli struct {
+		Anonymous `envprefix:"ANON_"`
+	}
+	w := &strings.Builder{}
+	p := mustNew(t, &cli, kong.Writers(w, w), kong.Exit(func(int) {}))
+	_, err := p.Parse([]string{"--help"})
+	require.NoError(t, err)
+	require.Contains(t, w.String(), "A flag ($ANON_FLAG).")
+}
+
 func TestCustomHelpFormatter(t *testing.T) {
 	var cli struct {
 		Flag string `env:"FLAG" help:"A flag."`
diff --git a/resolver_test.go b/resolver_test.go
index 2606ad4..fa83465 100644
--- a/resolver_test.go
+++ b/resolver_test.go
@@ -74,6 +74,39 @@ func TestEnvarsTag(t *testing.T) {
 	require.Equal(t, []int{5, 2, 9}, cli.Slice)
 }
 
+func TestEnvarsEnvPrefix(t *testing.T) {
+	type Anonymous struct {
+		Slice []int `env:"NUMBERS"`
+	}
+	var cli struct {
+		Anonymous `envprefix:"KONG_"`
+	}
+	parser, restoreEnv := newEnvParser(t, &cli, envMap{"KONG_NUMBERS": "1,2,3"})
+	defer restoreEnv()
+
+	_, err := parser.Parse([]string{})
+	require.NoError(t, err)
+	require.Equal(t, []int{1, 2, 3}, cli.Slice)
+}
+
+func TestEnvarsNestedEnvPrefix(t *testing.T) {
+	type NestedAnonymous struct {
+		String string `env:"STRING"`
+	}
+	type Anonymous struct {
+		NestedAnonymous `envprefix:"ANON_"`
+	}
+	var cli struct {
+		Anonymous `envprefix:"KONG_"`
+	}
+	parser, restoreEnv := newEnvParser(t, &cli, envMap{"KONG_ANON_STRING": "abc"})
+	defer restoreEnv()
+
+	_, err := parser.Parse([]string{})
+	require.NoError(t, err)
+	require.Equal(t, "abc", cli.String)
+}
+
 func TestEnvarsWithDefault(t *testing.T) {
 	var cli struct {
 		Flag string `env:"KONG_FLAG" default:"default"`
diff --git a/tag.go b/tag.go
index 84c068d..03b8ca6 100644
--- a/tag.go
+++ b/tag.go
@@ -32,6 +32,7 @@ type Tag struct {
 	Xor         []string
 	Vars        Vars
 	Prefix      string // Optional prefix on anonymous structs. All sub-flags will have this prefix.
+	EnvPrefix   string
 	Embed       bool
 	Aliases     []string
 	Negatable   bool
@@ -196,6 +197,7 @@ func hydrateTag(t *Tag, typeName string, isBool bool) error {
 		t.Xor = append(t.Xor, strings.FieldsFunc(xor, tagSplitFn)...)
 	}
 	t.Prefix = t.Get("prefix")
+	t.EnvPrefix = t.Get("envprefix")
 	t.Embed = t.Has("embed")
 	negatable := t.Has("negatable")
 	if negatable && !isBool {