github.com/kaydxh/golang@v0.0.131/go/flag/flag.go (about)

     1  /*
     2   *Copyright (c) 2022, kaydxh
     3   *
     4   *Permission is hereby granted, free of charge, to any person obtaining a copy
     5   *of this software and associated documentation files (the "Software"), to deal
     6   *in the Software without restriction, including without limitation the rights
     7   *to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     8   *copies of the Software, and to permit persons to whom the Software is
     9   *furnished to do so, subject to the following conditions:
    10   *
    11   *The above copyright notice and this permission notice shall be included in all
    12   *copies or substantial portions of the Software.
    13   *
    14   *THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    15   *IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    16   *FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    17   *AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    18   *LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    19   *OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    20   *SOFTWARE.
    21   */
    22  package flag
    23  
    24  import (
    25  	"bytes"
    26  	"fmt"
    27  	"io"
    28  	"strings"
    29  
    30  	"github.com/spf13/cobra"
    31  	"github.com/spf13/pflag"
    32  )
    33  
    34  const (
    35  	usageFmt = "Usage:\n  %s\n"
    36  )
    37  
    38  // NamedFlagSets stores named flag sets in the order of calling FlagSet.
    39  type NamedFlagSets struct {
    40  	// Order is an ordered list of flag set names.
    41  	Order []string
    42  	// FlagSets stores the flag sets by name.
    43  	FlagSets map[string]*pflag.FlagSet
    44  	// NormalizeNameFunc is the normalize function which used to initialize FlagSets created by NamedFlagSets.
    45  	NormalizeNameFunc func(f *pflag.FlagSet, name string) pflag.NormalizedName
    46  }
    47  
    48  // FlagSet returns the flag set with the given name and adds it to the
    49  // ordered name list if it is not in there yet.
    50  func (nfs *NamedFlagSets) FlagSet(name string) *pflag.FlagSet {
    51  	if nfs.FlagSets == nil {
    52  		nfs.FlagSets = map[string]*pflag.FlagSet{}
    53  	}
    54  	if _, ok := nfs.FlagSets[name]; !ok {
    55  		flagSet := pflag.NewFlagSet(name, pflag.ExitOnError)
    56  		flagSet.SetNormalizeFunc(pflag.CommandLine.GetNormalizeFunc())
    57  		if nfs.NormalizeNameFunc != nil {
    58  			flagSet.SetNormalizeFunc(nfs.NormalizeNameFunc)
    59  		}
    60  		nfs.FlagSets[name] = flagSet
    61  		nfs.Order = append(nfs.Order, name)
    62  	}
    63  	return nfs.FlagSets[name]
    64  }
    65  
    66  // PrintSections prints the given names flag sets in sections, with the maximal given column number.
    67  // If cols is zero, lines are not wrapped.
    68  func PrintSections(w io.Writer, fss NamedFlagSets, cols int) {
    69  	for _, name := range fss.Order {
    70  		fs := fss.FlagSets[name]
    71  		if !fs.HasFlags() {
    72  			continue
    73  		}
    74  
    75  		wideFS := pflag.NewFlagSet("", pflag.ExitOnError)
    76  		wideFS.AddFlagSet(fs)
    77  
    78  		var zzz string
    79  		if cols > 24 {
    80  			zzz = strings.Repeat("z", cols-24)
    81  			wideFS.Int(zzz, 0, strings.Repeat("z", cols-24))
    82  		}
    83  
    84  		var buf bytes.Buffer
    85  		fmt.Fprintf(&buf, "\n%s flags:\n\n%s", strings.ToUpper(name[:1])+name[1:], wideFS.FlagUsagesWrapped(cols))
    86  
    87  		if cols > 24 {
    88  			i := strings.Index(buf.String(), zzz)
    89  			lines := strings.Split(buf.String()[:i], "\n")
    90  			fmt.Fprint(w, strings.Join(lines[:len(lines)-1], "\n"))
    91  			fmt.Fprintln(w)
    92  		} else {
    93  			fmt.Fprint(w, buf.String())
    94  		}
    95  	}
    96  }
    97  
    98  // SetUsageAndHelpFunc set both usage and help function.
    99  // Print the flag sets we need instead of all of them.
   100  func SetUsageAndHelpFunc(cmd *cobra.Command, fss NamedFlagSets, cols int) {
   101  	cmd.SetUsageFunc(func(cmd *cobra.Command) error {
   102  		fmt.Fprintf(cmd.OutOrStderr(), usageFmt, cmd.UseLine())
   103  		PrintSections(cmd.OutOrStderr(), fss, cols)
   104  		return nil
   105  	})
   106  	cmd.SetHelpFunc(func(cmd *cobra.Command, args []string) {
   107  		fmt.Fprintf(cmd.OutOrStdout(), "%s\n\n"+usageFmt, cmd.Long, cmd.UseLine())
   108  		PrintSections(cmd.OutOrStdout(), fss, cols)
   109  	})
   110  }