dehub/cmd/dehub/dcmd/dcmd.go

171 lines
4.2 KiB
Go
Raw Normal View History

// 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
}
// 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) {
var title string
if cmd.parent == nil {
title = fmt.Sprintf("USAGE: %s [flags]", cmd.binary)
} else {
title = fmt.Sprintf("%s [%s flags]", cmd.name, cmd.name)
}
if subCmdTitle != "" {
title += " " + subCmdTitle
} else if len(cmd.subCmds) > 0 {
title += fmt.Sprint(" <sub-command> [sub-command flags]")
}
if cmd.parent == nil {
fmt.Printf("\n%s\n\n", title)
fmt.Print("### FLAGS ###\n\n")
} else {
cmd.parent.printUsageHead(title)
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)
}