utils/cobrautil/templates: add prefix based flags grouping

.golangci.yml

@@ -133,3 +133,7 @@ issues:
     - linters:
         - nosnakecase
       source: "func Example"
+
+    - linters:
+        - gomnd
+      source: "Priority:"

cmd/forwarder/root.go

@@ -37,7 +37,31 @@ func rootCommand() *cobra.Command {
 		},
 	}
 	commandGroups.Add(cmd)
-	templates.ActsAsRootCommand(cmd, nil, commandGroups...)
+
+	flagGroups := templates.FlagGroups{
+		{
+			Name:   "Flags",
+			Prefix: "",
+		},
+		{
+			Name:   "API server flags",
+			Prefix: "api",
+		},
+		{
+			Name:   "DNS flags",
+			Prefix: "dns",
+		},
+		{
+			Name:   "HTTP client flags",
+			Prefix: "http",
+		},
+		{
+			Name:   "Logging flags",
+			Prefix: "log",
+		},
+	}
+
+	templates.ActsAsRootCommand(cmd, nil, commandGroups, flagGroups)
 
 	// Add other commands
 	cmd.AddCommand(

utils/cobrautil/templates/flag_groups.go

@@ -0,0 +1,60 @@
+// Copyright 2023 Sauce Labs Inc. All rights reserved.
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+package templates
+
+import (
+	"sort"
+	"strings"
+
+	"github.com/spf13/pflag"
+)
+
+type FlagGroup struct {
+	Name   string
+	Prefix string
+
+	priority int
+}
+
+type FlagGroups []FlagGroup
+
+// processedFlagGroups returns a copy of the flag groups with the priority field set to the index of the group.
+// The returned groups are sorted by the length of the prefix in descending order.
+func processedFlagGroups(g FlagGroups) FlagGroups {
+	c := make(FlagGroups, len(g))
+	copy(c, g)
+
+	for i := range c {
+		c[i].priority = i
+	}
+
+	sort.Slice(c, func(i, j int) bool {
+		return len(c[i].Prefix) > len(c[j].Prefix)
+	})
+
+	return c
+}
+
+// splitFlagSet splits a flag set into multiple flag sets based on the prefix of the flag names.
+// If multiple groups match a flag, the flag is added to the first matching group.
+// The returned flag sets are ordered by the order of the groups.
+func (g FlagGroups) splitFlagSet(f *pflag.FlagSet) []*pflag.FlagSet {
+	var result []*pflag.FlagSet
+	for _, p := range g {
+		result = append(result, pflag.NewFlagSet(p.Name, pflag.ExitOnError))
+	}
+
+	f.VisitAll(func(f *pflag.Flag) {
+		for i := range g {
+			if strings.HasPrefix(f.Name, g[i].Prefix) {
+				result[i].AddFlag(f)
+				break
+			}
+		}
+	})
+	return result
+}

utils/cobrautil/templates/templater.go

@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"fmt"
 	"regexp"
+	"sort"
 	"strings"
 	"text/template"
 	"unicode"
@@ -35,15 +36,17 @@ type FlagExposer interface {
 	ExposeFlags(cmd *cobra.Command, flags ...string) FlagExposer
 }
 
-func ActsAsRootCommand(cmd *cobra.Command, filters []string, groups ...CommandGroup) FlagExposer {
+func ActsAsRootCommand(cmd *cobra.Command, filters []string, cg CommandGroups, fg FlagGroups) FlagExposer {
 	if cmd == nil {
 		panic("nil root command")
 	}
+
 	templater := &templater{
 		RootCmd:       cmd,
 		UsageTemplate: MainUsageTemplate(),
 		HelpTemplate:  MainHelpTemplate(),
-		CommandGroups: groups,
+		CommandGroups: cg,
+		FlagGroups:    processedFlagGroups(fg),
 		Filtered:      filters,
 	}
 	cmd.SetFlagErrorFunc(templater.FlagErrorFunc())
@@ -67,6 +70,7 @@ type templater struct {
 	HelpTemplate  string
 	RootCmd       *cobra.Command
 	CommandGroups
+	FlagGroups
 	Filtered []string
 }
 
@@ -229,13 +233,52 @@ func (t *templater) flagsUsages(f *flag.FlagSet) (string, error) {
 	flagBuf := new(bytes.Buffer)
 	printer := NewHelpFlagPrinter(flagBuf, t.RootCmd.Name(), DefaultWrapLimit)
 
-	f.VisitAll(func(flag *flag.Flag) {
-		if flag.Hidden {
-			return
-		}
-		printer.PrintHelpFlag(flag)
+	printFs := func() {
+		f.VisitAll(func(flag *flag.Flag) {
+			if flag.Hidden {
+				return
+			}
+			printer.PrintHelpFlag(flag)
+		})
+	}
+
+	g := t.FlagGroups
+
+	if len(g) == 0 {
+		fmt.Fprintf(flagBuf, "Flags:\n")
+		printFs()
+		return flagBuf.String(), nil
+	}
+
+	// Split flags into groups.
+	flagSet := g.splitFlagSet(f)
+
+	// Get permutation of groups sorted by priority.
+	perm := make([]int, len(g))
+	for i := range perm {
+		perm[i] = i
+	}
+	sort.Slice(perm, func(i, j int) bool {
+		return g[perm[i]].priority < g[perm[j]].priority
 	})
 
+	for _, i := range perm {
+		fs := flagSet[i]
+		if !fs.HasAvailableFlags() {
+			continue
+		}
+
+		if len(g[i].Name) > 0 {
+			fmt.Fprintf(flagBuf, "%s:\n", g[i].Name)
+		}
+		fs.VisitAll(func(flag *flag.Flag) {
+			if flag.Hidden {
+				return
+			}
+			printer.PrintHelpFlag(flag)
+		})
+	}
+
 	return flagBuf.String(), nil
 }
 

utils/cobrautil/templates/templates.go

@@ -48,8 +48,7 @@ const (
 {{end}}`
 
 	// SectionFlags is the help template section that displays the command's flags.
-	SectionFlags = `{{ if or $visibleFlags.HasFlags $explicitlyExposedFlags.HasFlags}}Options:
-{{ if $visibleFlags.HasFlags}}{{trimRight (flagsUsages $visibleFlags)}}{{end}}{{ if $explicitlyExposedFlags.HasFlags}}{{ if $visibleFlags.HasFlags}}
+	SectionFlags = `{{ if or $visibleFlags.HasFlags $explicitlyExposedFlags.HasFlags}}{{ if $visibleFlags.HasFlags}}{{trimRight (flagsUsages $visibleFlags)}}{{end}}{{ if $explicitlyExposedFlags.HasFlags}}{{ if $visibleFlags.HasFlags}}
 {{end}}{{trimRight (flagsUsages $explicitlyExposedFlags)}}{{end}}
 
 {{end}}`
@@ -85,7 +84,7 @@ func MainUsageTemplate() string {
 		SectionFlags,
 		SectionUsage,
 		SectionTipsHelp,
-		//SectionTipsGlobalOptions,
+		// SectionTipsGlobalOptions,
 	}
 	return strings.TrimRightFunc(strings.Join(sections, ""), unicode.IsSpace)
 }