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