set.go

// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package flag

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"strings"

	"github.com/posener/complete"
)

// Sets is a group of flag sets.
type Sets struct {
	// unionSet is the set that is the union of all other sets. This
	// has ALL flags defined on it and is the set that is parsed. But
	// we maintain the other list of sets so that we can generate proper help.
	unionSet *flag.FlagSet

	// flagSets is the list of sets that we have. We don't parse these
	// directly but use them for help generation and autocompletion.
	flagSets []*Set

	// completions is our set of autocompletion handlers. This is also
	// the union of all available flags similar to unionSet.
	completions complete.Flags
}

// NewSets creates a new flag sets.
func NewSets() *Sets {
	unionSet := flag.NewFlagSet("", flag.ContinueOnError)

	// Errors and usage are expected to be controlled externally by
	// checking on the result of Parse.
	unionSet.Usage = func() {}
	unionSet.SetOutput(io.Discard)

	return &Sets{
		unionSet:    unionSet,
		completions: complete.Flags{},
	}
}

// NewSet creates a new single flag set. A set should be created for
// any grouping of flags, for example "Common Options", "Auth Options", etc.
func (s *Sets) NewSet(name string) *Set {
	flagSet := NewSet(name)

	// The union and completions are pointers to our own values
	flagSet.unionSet = s.unionSet
	flagSet.completions = s.completions

	// Keep track of it for help generation
	s.flagSets = append(s.flagSets, flagSet)
	return flagSet
}

// GetSetFlags returns a slice of flags for a given set.
// If the requested set does not exist, this will return an empty slice.
func (s *Sets) GetSetFlags(setName string) []string {
	var setFlags []string
	for _, set := range s.flagSets {
		if set.name == setName {
			set.flagSet.VisitAll(func(f *flag.Flag) {
				setFlags = append(setFlags, fmt.Sprintf("-%s", f.Name))
			})
			return setFlags
		}
	}

	return setFlags
}

// Completions returns the completions for this flag set.
func (s *Sets) Completions() complete.Flags {
	return s.completions
}

// Parse parses the given flags, returning any errors.
func (s *Sets) Parse(args []string) error {
	return s.unionSet.Parse(args)
}

// Parsed reports whether the command-line flags have been parsed.
func (s *Sets) Parsed() bool {
	return s.unionSet.Parsed()
}

// Args returns the remaining args after parsing.
func (s *Sets) Args() []string {
	return s.unionSet.Args()
}

// Visit visits the flags in lexicographical order, calling fn for each. It
// visits only those flags that have been set.
func (s *Sets) Visit(fn func(*flag.Flag)) {
	s.unionSet.Visit(fn)
}

// Help builds custom help for this command, grouping by flag set.
func (s *Sets) Help() string {
	var out bytes.Buffer

	for _, set := range s.flagSets {
		printFlagTitle(&out, set.name+":")
		set.VisitAll(func(f *flag.Flag) {
			// Skip any hidden flags
			if v, ok := f.Value.(FlagVisibility); ok && v.Hidden() {
				return
			}
			printFlagDetail(&out, f)
		})
	}

	return strings.TrimRight(out.String(), "\n")
}

// VisitSets visits each set and performs action based on the passed in
// function.
func (s *Sets) VisitSets(fn func(name string, set *Set)) {
	for _, set := range s.flagSets {
		fn(set.name, set)
	}
}

// Set is a grouped wrapper around a real flag set and a grouped flag set.
type Set struct {
	name        string
	flagSet     *flag.FlagSet
	unionSet    *flag.FlagSet
	completions complete.Flags

	vars []*VarFlag
}

// NewSet creates a new flag set.
func NewSet(name string) *Set {
	return &Set{
		name:    name,
		flagSet: flag.NewFlagSet(name, flag.ContinueOnError),
	}
}

// Name returns the name of this flag set.
func (f *Set) Name() string {
	return f.name
}

func (f *Set) Visit(fn func(*flag.Flag)) {
	f.flagSet.Visit(fn)
}

func (f *Set) VisitAll(fn func(*flag.Flag)) {
	f.flagSet.VisitAll(fn)
}

func (f *Set) VisitVars(fn func(*VarFlag)) {
	for _, v := range f.vars {
		fn(v)
	}
}

// printFlagTitle prints a consistently-formatted title to the given writer.
func printFlagTitle(w io.Writer, s string) {
	fmt.Fprintf(w, "%s\n\n", s)
}

// printFlagDetail prints a single flag to the given writer.
func printFlagDetail(w io.Writer, f *flag.Flag) {
	// Check if the flag is hidden - do not print any flag detail or help output
	// if it is hidden.
	if h, ok := f.Value.(FlagVisibility); ok && h.Hidden() {
		return
	}

	// Check for a detailed example
	example := ""
	if t, ok := f.Value.(FlagExample); ok {
		example = t.Example()
	}

	if example != "" {
		fmt.Fprintf(w, "  -%s=<%s>\n", f.Name, example)
	} else {
		fmt.Fprintf(w, "  -%s\n", f.Name)
	}

	usage := reRemoveWhitespace.ReplaceAllString(f.Usage, " ")
	indented := wrapAtLengthWithPadding(usage, 6)
	fmt.Fprintf(w, "%s\n\n", indented)
}