blob: 72d42167d74f2aca02326d231f5b661f8cd6456f [file] [log] [blame]
// Copyright 2012 Marc-Antoine Ruel. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.
// Package subcommands permits a Go application to implement subcommands support
// similar to what is supported by the 'go' tool.
//
// The library is designed so that the test cases can run concurrently.
// Using global flags variables is discouraged to keep your program testable
// conccurently.
package subcommands
import (
"flag"
"fmt"
"io"
"os"
"sort"
"strings"
"sync"
"text/template"
"github.com/texttheater/golang-levenshtein/levenshtein"
)
// Application describes an application with subcommand support.
type Application interface {
// GetName returns the 'name' of the application.
GetName() string
// GetTitle returns a one-line title explaining the purpose of the tool.
GetTitle() string
// GetCommands returns the list of the subcommands that are supported.
GetCommands() []*Command
// GetOut is used for testing to allow parallel test case execution, should
// be normally os.Stdout.
GetOut() io.Writer
// GetErr is used for testing to allow parallel test case execution, should
// be normally os.Stderr.
GetErr() io.Writer
// GetEnvVars returns the map of EnvVarName -> EnvVarDefinition that this
// Application responds to.
GetEnvVars() map[string]EnvVarDefinition
}
// EnvVarDefinition describes an environment variable that this application
// responds to.
type EnvVarDefinition struct {
Advanced bool
ShortDesc string
Default string
}
// DefaultApplication implements all of Application interface's methods. An
// application should usually have a global instance of DefaultApplication and
// route main() to command_support.Run(app).
type DefaultApplication struct {
Name string
Title string
Commands []*Command
EnvVars map[string]EnvVarDefinition
}
// GetName implements interface Application.
func (a *DefaultApplication) GetName() string {
return a.Name
}
// GetTitle implements interface Application.
func (a *DefaultApplication) GetTitle() string {
return a.Title
}
// GetCommands implements interface Application.
func (a *DefaultApplication) GetCommands() []*Command {
return a.Commands
}
// GetOut implements interface Application.
func (a *DefaultApplication) GetOut() io.Writer {
return os.Stdout
}
// GetErr implements interface Application.
func (a *DefaultApplication) GetErr() io.Writer {
return os.Stderr
}
// GetEnvVars implements interface Application.
func (a *DefaultApplication) GetEnvVars() map[string]EnvVarDefinition {
return a.EnvVars
}
// Env is the mapping of resolved environment variables passed to
// CommandRun.Run.
type Env map[string]EnvVar
// EnvVar will document the value and existence of a given environment variable,
// as defined by Application.GetEnvVars. Value will be the value from the
// environment, or the Default value if it didn't exist. Exists will be true iff
// the value was present in the environment.
type EnvVar struct {
Value string
Exists bool
}
// CommandRun is an initialized object representing a subcommand that is ready
// to be executed.
type CommandRun interface {
// Run execute the actual command. When this function is called by
// command_support.Run(), the flags have already been parsed.
Run(a Application, args []string, env Env) int
// GetFlags returns the flags for this specific command.
//
// If this returns `nil`, then any additional arguments for this command will
// be passed unaltered as `args` to `Run()`. This is useful to delay command
// line parsing for implementing, for example, wrapper commands around other
// scripts.
GetFlags() *flag.FlagSet
}
// CommandRunBase implements GetFlags of CommandRun. It should be embedded in
// another struct that implements Run().
type CommandRunBase struct {
Flags flag.FlagSet
}
// GetFlags implements CommandRun.
func (c *CommandRunBase) GetFlags() *flag.FlagSet {
return &c.Flags
}
// Command describes a subcommand. It has one generator to generate a command
// object which is executable. The purpose of this design is to enable safe
// parallel execution of test cases.
type Command struct {
UsageLine string
ShortDesc string
LongDesc string
Advanced bool
CommandRun func() CommandRun
isSection bool
}
// Name returns the command's name: the first word in the usage line.
func (c *Command) Name() string {
name := c.UsageLine
i := strings.Index(name, " ")
if i >= 0 {
name = name[:i]
}
return name
}
// Section returns an un-runnable command that can act as a nice section
// heading for other commands.
func Section(name string) *Command {
return &Command{
ShortDesc: "\n\t" + name,
isSection: true,
}
}
// Usage prints out the general application Usage.
//
// This is primarily useful when using embedded commands. See sample-complex
// for an example.
//
// Beware that using the form "<tool> help -advanced <command>" will not
// propagate CmdHelp's help into the subcommand help Usage.
func Usage(out io.Writer, a Application, includeAdvanced bool) {
usageTemplate := `{{.Title}}
Usage: {{.Name}} [command] [arguments]
Commands:{{range .Commands}}
{{.Name | printf "%%-%ds"}} {{.ShortDesc}}{{end}}
{{if .EnvVars}}Environment Variables:{{range .EnvVars}}
{{.Name | printf "%%-%ds"}} {{.ShortDesc}}{{if .Default}} (Default: {{.Default | printf "%%q"}}){{end}}{{end}}
{{end}}
Use "{{.Name}} help [command]" for more information about a command.{{if .ShowAdvancedTip}}
Use "{{.Name}} help -advanced" to display all commands.{{end}}
`
widestCmd := 0
allCmds := a.GetCommands()
cmds := make([]*Command, 0, len(allCmds))
hasAdvanced := false
for _, c := range allCmds {
hasAdvanced = hasAdvanced || c.Advanced
if !c.Advanced || includeAdvanced {
// We need to include this command
if namLen := len(c.Name()); namLen > widestCmd {
widestCmd = namLen
}
cmds = append(cmds, c)
}
}
type envVarEntry struct {
Name string
ShortDesc string
Default string
}
widestEnvVar := 0
envVars := []envVarEntry(nil)
if envVarMap := a.GetEnvVars(); len(envVarMap) > 0 {
envVarKeys := make(sort.StringSlice, 0, len(envVarMap))
for k, v := range envVarMap {
if v.Advanced {
hasAdvanced = true
}
if !v.Advanced || includeAdvanced {
if keyLen := len(k); keyLen > widestEnvVar {
widestEnvVar = keyLen
}
envVarKeys = append(envVarKeys, k)
}
}
envVarKeys.Sort()
envVars = make([]envVarEntry, 0, len(envVarKeys))
for _, k := range envVarKeys {
v := envVarMap[k]
envVars = append(envVars, envVarEntry{k, v.ShortDesc, v.Default})
}
}
data := map[string]interface{}{
"Title": a.GetTitle(),
"Name": a.GetName(),
"Commands": cmds,
"EnvVars": envVars,
"ShowAdvancedTip": (hasAdvanced && !includeAdvanced),
}
tmpl(out, fmt.Sprintf(usageTemplate, widestCmd, widestEnvVar), data)
}
// getCommandUsageHandler returns a flag.Usage compatible function.
func getCommandUsageHandler(out io.Writer, a Application, c *Command, r CommandRun, helpUsed *bool) func() {
return func() {
helpTemplate := "{{.Cmd.LongDesc | trim | wrapWithLines}}usage: {{.App.GetName}} {{.Cmd.UsageLine}}\n"
dict := struct {
App Application
Cmd *Command
}{a, c}
tmpl(out, helpTemplate, dict)
if f := r.GetFlags(); f != nil {
f.PrintDefaults()
}
*helpUsed = true
}
}
// Initializes the flags for a specific CommandRun.
func initCommand(a Application, c *Command, r CommandRun, out io.Writer, helpUsed *bool) (hasFlags bool) {
f := r.GetFlags()
if f != nil {
if f.Usage == nil {
f.Usage = getCommandUsageHandler(out, a, c, r, helpUsed)
}
f.SetOutput(out)
f.Init(c.Name(), flag.ContinueOnError)
}
return f != nil
}
// FindCommand finds a Command by name and returns it if found.
func FindCommand(a Application, name string) *Command {
for _, c := range a.GetCommands() {
if c.Name() == name {
return c
}
}
return nil
}
// FindNearestCommand heuristically finds a Command the user wanted to type but
// failed to type correctly.
func FindNearestCommand(a Application, name string) *Command {
commands := map[string]*Command{}
for _, c := range a.GetCommands() {
if !c.isSection {
commands[c.Name()] = c
}
}
if c, ok := commands[name]; ok {
return c
}
// Search for unique prefix.
withPrefix := []*Command{}
for n, c := range commands {
if strings.HasPrefix(n, name) {
withPrefix = append(withPrefix, c)
}
}
if len(withPrefix) == 1 {
return withPrefix[0]
}
// Search for case insensitivity.
withPrefix = []*Command{}
lowName := strings.ToLower(name)
for n, c := range commands {
if strings.HasPrefix(strings.ToLower(n), lowName) {
withPrefix = append(withPrefix, c)
}
}
if len(withPrefix) == 1 {
return withPrefix[0]
}
// Calculate the levenshtein distance and take the closest one.
closestD := 1000
var closestC *Command
secondD := 1000
for n, c := range commands {
dist := levenshtein.DistanceForStrings([]rune(n), []rune(name), levenshtein.DefaultOptions)
if dist < closestD {
secondD = closestD
closestD = dist
closestC = c
} else if dist < secondD {
secondD = dist
}
}
if closestD > 3 {
// Not similar enough. Don't be a fool and run a random command.
return nil
}
if (secondD - closestD) < 3 {
// Too ambiguous.
return nil
}
return closestC
}
// Run runs the application, scheduling the subcommand. This is the main entry
// point of the library.
//
// Unit tests should call this function directly with args provided so this is
// concurrent safe.
//
// It is safer to use a base class embedding CommandRunBase that is then
// embedded by each CommandRun implementation to define flags available for
// all commands.
func Run(a Application, args []string) int {
// Process general flags first, mainly for -help.
helpUsed := false
if args == nil {
// Do not parse during unit tests because flag.commandLine.errorHandling == ExitOnError. :(
args, helpUsed = parseGeneral(a)
}
if len(args) < 1 {
// Need a command.
Usage(a.GetErr(), a, false)
return 2
}
if c := FindNearestCommand(a, args[0]); c != nil {
// Initialize the flags.
r := c.CommandRun()
hasFlags := initCommand(a, c, r, a.GetErr(), &helpUsed)
var cmdArgs []string
if hasFlags {
if err := r.GetFlags().Parse(args[1:]); err != nil {
return 2
}
if helpUsed {
return 0
}
cmdArgs = r.GetFlags().Args()
} else {
cmdArgs = args[1:]
}
envVars := a.GetEnvVars()
envMap := make(map[string]EnvVar, len(envVars))
for k, v := range envVars {
val, ok := os.LookupEnv(k)
if !ok {
val = v.Default
}
envMap[k] = EnvVar{val, ok}
}
return r.Run(a, cmdArgs, envMap)
}
fmt.Fprintf(a.GetErr(), "%s: unknown command %#q\n\nRun '%s help' for usage.\n", a.GetName(), args[0], a.GetName())
return 2
}
var mu sync.Mutex
// parseGeneral parses the general flag in a way that is safe in unit tests
// even when using t.Parallel()
func parseGeneral(a Application) ([]string, bool) {
mu.Lock()
prev := flag.Usage
defer func() {
mu.Unlock()
flag.Usage = prev
}()
helpUsed := false
flag.Usage = func() {
Usage(a.GetErr(), a, false)
helpUsed = true
}
// This may call flag.Usage():
flag.Parse()
return flag.Args(), helpUsed
}
// tmpl executes the given template text on data, writing the result to w.
func tmpl(w io.Writer, text string, data interface{}) {
t := template.New("top")
t.Funcs(template.FuncMap{"trim": strings.TrimSpace, "wrapWithLines": wrapWithLines})
template.Must(t.Parse(text))
if err := t.Execute(w, data); err != nil {
panic(fmt.Sprintf("Failed to execute template: %s", err))
}
}
func wrapWithLines(s string) string {
if s == "" {
return s
}
return s + "\n\n"
}
// CmdHelp defines the help command. It should be included in your application's
// Commands list.
//
// It is not added automatically but it will be run automatically if added.
var CmdHelp = &Command{
UsageLine: "help [<command>|-advanced]",
ShortDesc: "prints help about a command",
LongDesc: "Prints an overview of every command or information about a specific command.\nPass -advanced to see help for advanced commands.",
CommandRun: func() CommandRun {
ret := &helpRun{}
ret.Flags.BoolVar(&ret.advanced, "advanced", false, "show advanced commands")
return ret
},
}
type helpRun struct {
CommandRunBase
advanced bool
}
func (c *helpRun) Run(a Application, args []string, env Env) int {
if len(args) == 0 {
Usage(a.GetOut(), a, c.advanced)
return 0
}
if len(args) != 1 {
fmt.Fprintf(a.GetErr(), "%s: Too many arguments given\n\nRun '%s help' for usage.\n", a.GetName(), a.GetName())
return 2
}
// Redirects all output to Out.
var helpUsed bool
if cmd := FindNearestCommand(a, args[0]); cmd != nil {
// Initialize the flags.
r := cmd.CommandRun()
if initCommand(a, cmd, r, a.GetErr(), &helpUsed) {
r.GetFlags().Usage()
} else {
getCommandUsageHandler(a.GetErr(), a, cmd, r, &helpUsed)()
}
return 0
}
fmt.Fprintf(a.GetErr(), "%s: unknown command %#q\n\nRun '%s help' for usage.\n", a.GetName(), args[0], a.GetName())
return 2
}