dehub/cmd/dehub/dcmd/dcmd.go

// 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> [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)
}