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 }