add namedflags for grouping together flags

namedflagsets.go

@@ -0,0 +1,169 @@
+package cobrautil
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"math/rand"
+	"strings"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+)
+
+// MarkFlagsHidden is a convenient way to mark flags as hidden in bulk.
+func MarkFlagsHidden(flags *pflag.FlagSet, names ...string) error {
+	for _, name := range names {
+		if err := flags.MarkHidden(name); err != nil {
+			return fmt.Errorf("failed to mark flag as hidden: %w", err)
+		}
+	}
+	return nil
+}
+
+// NewNamedFlagSets creates a new NamedFlagSets and registers it with the
+// provided command.
+func NewNamedFlagSets(cmd *cobra.Command) *NamedFlagSets {
+	nfs := &NamedFlagSets{}
+	nfs.SetUsageTemplate(cmd)
+	return nfs
+}
+
+// SetUsageTemplate overrides the cobra usage template to include the
+// named flags sections.
+func (nfs *NamedFlagSets) SetUsageTemplate(cmd *cobra.Command) {
+	cobra.AddTemplateFunc(nfs.templateFuncName(), nfs.templateFunc)
+	cmd.SetUsageTemplate(nfs.usageTemplate())
+}
+
+// NamedFlagSets stores named flag sets in the order of calling FlagSet.
+//
+// This type is largely adapted from [k8s.io/component-base/cli/flag], but
+// modified to be less brittle by integrating with cobra's templates rather
+// than entirely overriding UsageFuncs and HelpFuncs.
+type NamedFlagSets struct {
+	// Order is an ordered list of flag set names.
+	Order []string
+
+	// FlagSets stores the flag sets by name.
+	FlagSets map[string]*pflag.FlagSet
+
+	// NormalizeNameFunc is the normalize function which used to initialize
+	// FlagSets created by NamedFlagSets.
+	NormalizeNameFunc func(f *pflag.FlagSet, name string) pflag.NormalizedName
+
+	uniqueID int
+}
+
+// templateFuncName generates a random template name so that template, which
+// has to be registered globally, isn't overridden by any other NamedFlagSets.
+func (nfs *NamedFlagSets) templateFuncName() string {
+	if nfs.uniqueID == 0 {
+		nfs.uniqueID = rand.Int()
+	}
+	return fmt.Sprintf("namedFlagSets%d", nfs.uniqueID)
+}
+
+func hasVisibleFlags(flags *pflag.FlagSet) bool {
+	var found bool
+	flags.VisitAll(func(flag *pflag.Flag) {
+		if !flag.Hidden {
+			found = true
+		}
+	})
+	return found
+}
+
+func (nfs *NamedFlagSets) AddFlagSets(cmd *cobra.Command) {
+	for _, name := range nfs.Order {
+		cmd.Flags().AddFlagSet(nfs.FlagSet(name))
+	}
+}
+
+// FlagSet returns the flag set with the given name and adds it to the
+// ordered name list if it is not in there yet.
+func (nfs *NamedFlagSets) FlagSet(name string) *pflag.FlagSet {
+	if nfs.FlagSets == nil {
+		nfs.FlagSets = map[string]*pflag.FlagSet{}
+	}
+	if _, ok := nfs.FlagSets[name]; !ok {
+		flagSet := pflag.NewFlagSet(name, pflag.ExitOnError)
+		flagSet.SetNormalizeFunc(pflag.CommandLine.GetNormalizeFunc())
+		if nfs.NormalizeNameFunc != nil {
+			flagSet.SetNormalizeFunc(nfs.NormalizeNameFunc)
+		}
+		nfs.FlagSets[name] = flagSet
+		nfs.Order = append(nfs.Order, name)
+	}
+	return nfs.FlagSets[name]
+}
+
+// printSections prints the given names flag sets in sections, with the maximal
+// given column number.
+//
+// If cols is zero, lines are not wrapped.
+func (nfs *NamedFlagSets) printSections(w io.Writer, cols int) {
+	for _, name := range nfs.Order {
+		fs := nfs.FlagSets[name]
+		if !hasVisibleFlags(fs) {
+			continue
+		}
+
+		wideFS := pflag.NewFlagSet("", pflag.ExitOnError)
+		wideFS.AddFlagSet(fs)
+
+		var zzz string
+		if cols > 24 {
+			zzz = strings.Repeat("z", cols-24)
+			wideFS.Int(zzz, 0, strings.Repeat("z", cols-24))
+		}
+
+		var buf bytes.Buffer
+		fmt.Fprintf(&buf, "\n%s Flags:\n%s", name, wideFS.FlagUsagesWrapped(cols))
+
+		if cols > 24 {
+			i := strings.Index(buf.String(), zzz)
+			lines := strings.Split(buf.String()[:i], "\n")
+			fmt.Fprint(w, strings.Join(lines[:len(lines)-1], "\n"))
+			fmt.Fprintln(w)
+		} else {
+			fmt.Fprint(w, buf.String())
+		}
+	}
+}
+
+func (nfs *NamedFlagSets) templateFunc() string {
+	var b strings.Builder
+	nfs.printSections(&b, 0)
+	return b.String()
+}
+
+func (nfs *NamedFlagSets) usageTemplate() string {
+	return `Usage:{{if .Runnable}}
+  {{.UseLine}}{{end}}{{if .HasAvailableSubCommands}}
+  {{.CommandPath}} [command]{{end}}{{if gt (len .Aliases) 0}}
+
+Aliases:
+  {{.NameAndAliases}}{{end}}{{if .HasExample}}
+
+Examples:
+{{.Example}}{{end}}{{if .HasAvailableSubCommands}}{{$cmds := .Commands}}{{if eq (len .Groups) 0}}
+
+Available Commands:{{range $cmds}}{{if (or .IsAvailableCommand (eq .Name "help"))}}
+  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{else}}{{range $group := .Groups}}
+
+{{.Title}}{{range $cmds}}{{if (and (eq .GroupID $group.ID) (or .IsAvailableCommand (eq .Name "help")))}}
+  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{if not .AllChildCommandsHaveGroup}}
+
+Additional Commands:{{range $cmds}}{{if (and (eq .GroupID "") (or .IsAvailableCommand (eq .Name "help")))}}
+  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{end}}{{end}}{{if .HasAvailableLocalFlags}}
+{{` + nfs.templateFuncName() + `}}
+Global Flags:
+{{.InheritedFlags.FlagUsages | trimTrailingWhitespaces}}{{end}}{{if .HasHelpSubCommands}}
+
+Additional help topics:{{range .Commands}}{{if .IsAdditionalHelpTopicCommand}}
+  {{rpad .CommandPath .CommandPathPadding}} {{.Short}}{{end}}{{end}}{{end}}{{if .HasAvailableSubCommands}}
+
+Use "{{.CommandPath}} [command] --help" for more information about a command.{{end}}
+`
+}