// Package dcmd implements command and sub-command parsing and runtime // management. It wraps the stdlib flag package as well, to incorporate // configuration into the mix. package dcmd import ( "context" "errors" "flag" "fmt" "os" "sort" "strings" ) func exitErr(err error) { fmt.Fprintf(os.Stderr, "exiting: %v\n", err) os.Stderr.Sync() os.Stdout.Sync() os.Exit(1) } type subCmd struct { name, descr string run func(context.Context, *Cmd) } // Cmd wraps a flag.FlagSet instance to provide extra functionality that dehub // wants, specifically around sub-command support. type Cmd struct { flagSet *flag.FlagSet binary string // only gets set on root Cmd, during Run subCmds []subCmd // these fields get set by the parent Cmd, if this is a sub-command. name string args []string parent *Cmd } // New initializes and returns an empty Cmd instance. func New() *Cmd { return &Cmd{} } func (cmd *Cmd) getFlagSet() *flag.FlagSet { if cmd.flagSet == nil { cmd.flagSet = flag.NewFlagSet(cmd.name, flag.ContinueOnError) } return cmd.flagSet } func (cmd *Cmd) numFlags() int { var n int cmd.getFlagSet().VisitAll(func(*flag.Flag) { n++ }) return n } // FlagSet returns a flag.Cmd instance on which parameter creation methods can // be called, e.g. String(...) or Int(...). func (cmd *Cmd) FlagSet() *flag.FlagSet { return cmd.getFlagSet() } // SubCmd registers a sub-command of this Cmd. // // A new Cmd will be instantiated when this sub-command is picked on the // command-line during this Cmd's Run method. The Context returned from that Run // and the new Cmd will be passed into the callback given here. The sub-command // should then be performed in the same manner as this Cmd is performed // (including setting flags, adding sub-sub-commands, etc...) func (cmd *Cmd) SubCmd(name, descr string, run func(context.Context, *Cmd)) { cmd.subCmds = append(cmd.subCmds, subCmd{ name: name, descr: descr, run: run, }) // it's not the most efficient to do this here, but it is the easiest sort.Slice(cmd.subCmds, func(i, j int) bool { return cmd.subCmds[i].name < cmd.subCmds[j].name }) } func (cmd *Cmd) printUsageHead(subCmdTitle string) { hasFlags := cmd.numFlags() > 0 var title string if cmd.parent == nil { title = fmt.Sprintf("USAGE: %s", cmd.binary) if hasFlags { title += " [flags]" } } else { title = fmt.Sprintf("%s", cmd.name) if hasFlags { title += fmt.Sprintf(" [%s flags]", cmd.name) } } if subCmdTitle != "" { title += " " + subCmdTitle } else if len(cmd.subCmds) > 0 { title += fmt.Sprint(" [sub-command flags]") } if cmd.parent == nil { fmt.Printf("\n%s\n\n", title) } else { cmd.parent.printUsageHead(title) } if hasFlags { if cmd.parent == nil { fmt.Print("### FLAGS ###\n\n") } else { fmt.Printf("### %s FLAGS ###\n\n", strings.ToUpper(cmd.name)) } cmd.getFlagSet().PrintDefaults() fmt.Print("\n") } } // Run performs the comand. It starts by parsing all flags in the Cmd's FlagSet, // and possibly exiting with a usage message if appropriate. It will then // perform the given body callback, and then perform any sub-commands (if // selected). // // The context returned from the callback will be passed into the callback // (given to SubCmd) of any sub-commands which are run, and so on. func (cmd *Cmd) Run(body func() (context.Context, error)) { args := cmd.args if cmd.parent == nil { cmd.binary, args = os.Args[0], os.Args[1:] } fs := cmd.getFlagSet() fs.Usage = func() { cmd.printUsageHead("") if len(cmd.subCmds) == 0 { return } fmt.Printf("### SUB-COMMANDS ###\n\n") for _, subCmd := range cmd.subCmds { fmt.Printf("\t%s : %s\n", subCmd.name, subCmd.descr) } fmt.Println("") } if err := fs.Parse(args); err != nil { exitErr(err) return } ctx, err := body() if err != nil { exitErr(err) } // body has run, now do sub-command (if there is one) subArgs := fs.Args() if len(cmd.subCmds) == 0 { return } else if len(subArgs) == 0 && len(cmd.subCmds) > 0 { fs.Usage() exitErr(errors.New("no sub-command selected")) } // now find that sub-command subCmdName := strings.ToLower(subArgs[0]) var subCmd subCmd var subCmdOk bool for _, subCmd = range cmd.subCmds { if subCmdOk = subCmd.name == subCmdName; subCmdOk { break } } if !subCmdOk { fs.Usage() exitErr(fmt.Errorf("unknown command %q", subCmdName)) } subCmdCmd := New() subCmdCmd.name = subCmd.name subCmdCmd.args = subArgs[1:] subCmdCmd.parent = cmd subCmd.run(ctx, subCmdCmd) }