From 7a635254200602a6de251035c6de8a9259abd8f1 Mon Sep 17 00:00:00 2001
From: Alec Thomas <alec@swapoff.org>
Date: Fri, 11 Mar 2022 22:51:12 +1100
Subject: [PATCH] feat: AutoGroup option

---
 help_test.go | 75 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 options.go   | 15 +++++++++++
 2 files changed, 90 insertions(+)

diff --git a/help_test.go b/help_test.go
index e0e9043..2018b0f 100644
--- a/help_test.go
+++ b/help_test.go
@@ -494,6 +494,81 @@ func TestCustomValueFormatter(t *testing.T) {
 	require.Contains(t, w.String(), "A flag.")
 }
 
+func TestAutoGroup(t *testing.T) {
+	var cli struct {
+		GroupedAString string `help:"A string flag grouped in A."`
+		FreeString     string `help:"A non grouped string flag."`
+		GroupedBString string `help:"A string flag grouped in B."`
+		FreeBool       bool   `help:"A non grouped bool flag."`
+		GroupedABool   bool   `help:"A bool flag grouped in A."`
+
+		One struct {
+			Flag string `help:"Nested flag."`
+			// Group is inherited from the parent command
+			Thing struct {
+				Arg string `arg help:"argument"`
+			} `cmd help:"subcommand thing"`
+			Other struct {
+				Other string `arg help:"other arg"`
+			} `arg help:"subcommand other"`
+			// ... but a subcommand can override it
+			Stuff struct {
+				Stuff string `arg help:"argument"`
+			} `arg help:"subcommand stuff"`
+		} `cmd help:"A subcommand grouped in A."`
+
+		Two struct {
+			Grouped1String  string `help:"A string flag grouped in 1."`
+			AFreeString     string `help:"A non grouped string flag."`
+			Grouped2String  string `help:"A string flag grouped in 2."`
+			AGroupedAString bool   `help:"A string flag grouped in A."`
+			Grouped1Bool    bool   `help:"A bool flag grouped in 1."`
+		} `cmd help:"A non grouped subcommand."`
+
+		Four struct {
+			Flag string `help:"Nested flag."`
+		} `cmd help:"Another subcommand grouped in B."`
+
+		Three struct {
+			Flag string `help:"Nested flag."`
+		} `cmd help:"Another subcommand grouped in A."`
+	}
+	w := bytes.NewBuffer(nil)
+	app := mustNew(t, &cli,
+		kong.Writers(w, w),
+		kong.Exit(func(int) {}),
+		kong.AutoGroup(func(parent kong.Visitable, flag *kong.Flag) *kong.Group {
+			if node, ok := parent.(*kong.Node); ok {
+				return &kong.Group{
+					Key:   node.Name,
+					Title: strings.Title(node.Name) + " flags:",
+				}
+			}
+			return nil
+		}),
+	)
+	_, _ = app.Parse([]string{"--help", "two"})
+	require.Equal(t, `Usage: test two
+
+A non grouped subcommand.
+
+Flags:
+  -h, --help                       Show context-sensitive help.
+      --grouped-a-string=STRING    A string flag grouped in A.
+      --free-string=STRING         A non grouped string flag.
+      --grouped-b-string=STRING    A string flag grouped in B.
+      --free-bool                  A non grouped bool flag.
+      --grouped-a-bool             A bool flag grouped in A.
+
+Two flags:
+  --grouped-1-string=STRING    A string flag grouped in 1.
+  --a-free-string=STRING       A non grouped string flag.
+  --grouped-2-string=STRING    A string flag grouped in 2.
+  --a-grouped-a-string         A string flag grouped in A.
+  --grouped-1-bool             A bool flag grouped in 1.
+`, w.String())
+}
+
 func TestHelpGrouping(t *testing.T) {
 	var cli struct {
 		GroupedAString string `help:"A string flag grouped in A." group:"Group A"`
diff --git a/options.go b/options.go
index 4c240ed..7ddb4e3 100644
--- a/options.go
+++ b/options.go
@@ -240,6 +240,21 @@ func ConfigureHelp(options HelpOptions) Option {
 	})
 }
 
+// AutoGroup automatically assigns groups to flags.
+func AutoGroup(format func(parent Visitable, flag *Flag) *Group) Option {
+	return PostBuild(func(kong *Kong) error {
+		parents := []Visitable{kong.Model}
+		return Visit(kong.Model, func(node Visitable, next Next) error {
+			if flag, ok := node.(*Flag); ok && flag.Group == nil {
+				flag.Group = format(parents[len(parents)-1], flag)
+			}
+			parents = append(parents, node)
+			defer func() { parents = parents[:len(parents)-1] }()
+			return next(nil)
+		})
+	})
+}
+
 // Groups associates `group` field tags with group metadata.
 //
 // This option is used to simplify Kong tags while providing