github.com/grafana/pyroscope@v1.18.0/pkg/cfg/flag.go (about)

     1  package cfg
     2  
     3  import (
     4  	"flag"
     5  	"fmt"
     6  	"sort"
     7  	"strings"
     8  
     9  	"github.com/grafana/dskit/flagext"
    10  	"github.com/pkg/errors"
    11  	"golang.org/x/text/cases"
    12  	"golang.org/x/text/language"
    13  )
    14  
    15  // Defaults registers flags to the flagSet using dst as the flagext.Registerer
    16  func Defaults(fs *flag.FlagSet) Source {
    17  	return func(dst Cloneable) error {
    18  		r, ok := dst.(flagext.Registerer)
    19  		if !ok {
    20  			return errors.New("dst does not satisfy flagext.Registerer")
    21  		}
    22  
    23  		// already sets the defaults on r
    24  		r.RegisterFlags(fs)
    25  		return nil
    26  	}
    27  }
    28  
    29  // Flags parses the flag from the command line, setting only user-supplied
    30  // values on the flagext.Registerer passed to Defaults()
    31  func Flags(args []string, fs *flag.FlagSet) Source {
    32  	flag.Usage = categorizedUsage(fs)
    33  	return dFlags(fs, args)
    34  }
    35  
    36  // dFlags parses the flagset, applying all values set on the slice
    37  func dFlags(fs *flag.FlagSet, args []string) Source {
    38  	return func(dst Cloneable) error {
    39  		// parse the final flagset
    40  		return fs.Parse(args)
    41  	}
    42  }
    43  
    44  func categorizedUsage(fs *flag.FlagSet) func() {
    45  	categories := make(map[string][]string)
    46  	return func() {
    47  		if fs.Name() == "" {
    48  			fmt.Fprintf(fs.Output(), "Usage:\n")
    49  		} else {
    50  			fmt.Fprintf(fs.Output(), "Usage of %s:\n", fs.Name())
    51  		}
    52  
    53  		fs.VisitAll(func(f *flag.Flag) {
    54  			id := ""
    55  			if strings.Contains(f.Name, ".") {
    56  				id = strings.Split(f.Name, ".")[0]
    57  			}
    58  
    59  			kind, usage := flag.UnquoteUsage(f)
    60  			if kind != "" {
    61  				kind = " " + kind
    62  			}
    63  			def := f.DefValue
    64  			if def != "" {
    65  				def = fmt.Sprintf(" (default %s)", def)
    66  			}
    67  			categories[id] = append(categories[id], fmt.Sprintf("   -%s%s:\n      %s%s", f.Name, kind, usage, def))
    68  		})
    69  
    70  		for name, flags := range categories {
    71  			if len(flags) == 1 {
    72  				categories[""] = append(categories[""], flags[0])
    73  				delete(categories, name)
    74  			}
    75  		}
    76  
    77  		for name := range categories {
    78  			sort.Strings(categories[name])
    79  		}
    80  
    81  		for _, u := range categories[""] {
    82  			fmt.Fprintln(fs.Output(), u)
    83  		}
    84  		fmt.Fprintln(fs.Output())
    85  
    86  		keys := make([]string, 0, len(categories))
    87  		for k := range categories {
    88  			keys = append(keys, k)
    89  		}
    90  		sort.Strings(keys)
    91  
    92  		for _, name := range keys {
    93  			if name == "" {
    94  				continue
    95  			}
    96  			fmt.Fprintf(fs.Output(), " %s:\n", cases.Title(language.English).String(name))
    97  			for _, u := range categories[name] {
    98  				fmt.Fprintln(fs.Output(), u)
    99  			}
   100  			fmt.Fprintln(fs.Output())
   101  		}
   102  	}
   103  }