Skip to content

Instantly share code, notes, and snippets.

@progrium
Created Sep 2, 2021
Embed
What would you like to do?
350 line reimplementation of cobra, simplified and with no dependencies
package cli
import (
"bytes"
"context"
"flag"
"fmt"
"os"
"reflect"
"strings"
"text/template"
"unicode"
)
var templateFuncs = template.FuncMap{
"trim": strings.TrimSpace,
"trimRight": trimRightSpace,
"rpad": rpad,
}
var HelpTemplate = `Usage:{{if .Runnable}}
{{.UseLine}}{{end}}{{if .HasAvailableSubCommands}}
{{.CommandPath}} [command]{{end}}{{if gt (len .Aliases) 0}}
Aliases:
{{.NameAndAliases}}{{end}}{{if .HasExample}}
Examples:
{{.Example}}{{end}}{{if .HasAvailableSubCommands}}
Available Commands:{{range .Commands}}{{if (or .IsAvailableCommand (eq .Name "help"))}}
{{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{if .HasAvailableLocalFlags}}
Flags:
{{.FlagUsages | trimRight }}{{end}}{{if .HasAvailableSubCommands}}
Use "{{.CommandPath}} [command] -help" for more information about a command.{{end}}
`
type Command struct {
// Use is the one-line usage message.
// Recommended syntax is as follow:
// [ ] identifies an optional argument. Arguments that are not enclosed in brackets are required.
// ... indicates that you can specify multiple values for the previous argument.
// | indicates mutually exclusive information. You can use the argument to the left of the separator or the
// argument to the right of the separator. You cannot use both arguments in a single use of the command.
// { } delimits a set of mutually exclusive arguments when one of the arguments is required. If the arguments are
// optional, they are enclosed in brackets ([ ]).
// Example: add [-F file | -D dir]... [-f format] <profile>
Usage string
// Short is the short description shown in the 'help' output.
Short string
// Long is the long message shown in the 'help <this-command>' output.
Long string
// Hidden defines, if this command is hidden and should NOT show up in the list of available commands.
Hidden bool
// Aliases is an array of aliases that can be used instead of the first word in Use.
Aliases []string
// Example is examples of how to use the command.
Example string
// Annotations are key/value pairs that can be used by applications to identify or
// group commands.
Annotations map[string]interface{}
// Version defines the version for this command. If this value is non-empty and the command does not
// define a "version" flag, a "version" boolean flag will be added to the command and, if specified,
// will print content of the "Version" variable. A shorthand "v" flag will also be added if the
// command does not define one.
Version string
Args PositionalArgs
Run func(ctx context.Context, args []string)
commands []*Command
parent *Command
flags *flag.FlagSet
}
func (c *Command) Flags() *flag.FlagSet {
if c.flags == nil {
c.flags = flag.NewFlagSet(c.Name(), flag.ContinueOnError)
var null bytes.Buffer
c.flags.SetOutput(&null)
}
return c.flags
}
func (c *Command) AddCommand(sub *Command) {
if sub == c {
panic("command can't be a child of itself")
}
sub.parent = c
c.commands = append(c.commands, sub)
}
// CommandPath returns the full path to this command.
func (c *Command) CommandPath() string {
if c.parent != nil {
return c.parent.CommandPath() + " " + c.Name()
}
return c.Name()
}
// UseLine puts out the full usage for a given command (including parents).
func (c *Command) UseLine() string {
if c.parent != nil {
return c.parent.CommandPath() + " " + c.Usage
} else {
return c.Usage
}
}
func (c *Command) Name() string {
name := c.Usage
i := strings.Index(name, " ")
if i >= 0 {
name = name[:i]
}
return name
}
func (c *Command) Find(args []string) (cmd *Command, n int) {
cmd = c
if len(args) == 0 {
return
}
var arg string
for n, arg = range args {
if cc := cmd.findSub(arg); cc != nil {
cmd = cc
} else {
return
}
}
n += 1
return
}
func (c *Command) findSub(name string) *Command {
for _, cmd := range c.commands {
if cmd.Name() == name || hasAlias(cmd, name) {
return cmd
}
}
return nil
}
func hasAlias(cmd *Command, name string) bool {
for _, a := range cmd.Aliases {
if a == name {
return true
}
}
return false
}
type CommandHelp struct {
*Command
}
func (c *CommandHelp) Runnable() bool {
return c.Run != nil
}
func (c *CommandHelp) IsAvailableCommand() bool {
if c.Hidden {
return false
}
if c.Runnable() || c.HasAvailableSubCommands() {
return true
}
return false
}
func (c *CommandHelp) HasAvailableSubCommands() bool {
for _, sub := range c.commands {
if (&CommandHelp{sub}).IsAvailableCommand() {
return true
}
}
return false
}
func (c *CommandHelp) NameAndAliases() string {
return strings.Join(append([]string{c.Name()}, c.Aliases...), ", ")
}
func (c *CommandHelp) HasExample() bool {
return len(c.Example) > 0
}
func (c *CommandHelp) Commands() (cmds []*CommandHelp) {
for _, cmd := range c.commands {
cmds = append(cmds, &CommandHelp{cmd})
}
return
}
func (c *CommandHelp) NamePadding() int {
return 16
}
func (c *CommandHelp) HasAvailableLocalFlags() bool {
n := 0
c.Flags().VisitAll(func(f *flag.Flag) {
n++
})
return n > 0
}
func (c *CommandHelp) FlagUsages() string {
var sb strings.Builder
c.Flags().VisitAll(func(f *flag.Flag) {
fmt.Fprintf(&sb, " -%s", f.Name) // Two spaces before -; see next two comments.
name, usage := flag.UnquoteUsage(f)
if len(name) > 0 {
sb.WriteString(" ")
sb.WriteString(name)
}
// Boolean flags of one ASCII letter are so common we
// treat them specially, putting their usage on the same line.
if sb.Len() <= 4 { // space, space, '-', 'x'.
sb.WriteString("\t")
} else {
// Four spaces before the tab triggers good alignment
// for both 4- and 8-space tab stops.
sb.WriteString("\n \t")
}
sb.WriteString(strings.ReplaceAll(usage, "\n", "\n \t"))
f.Usage = ""
if !isZeroValue(f, f.DefValue) {
typ, _ := flag.UnquoteUsage(f)
if typ == "string" {
// put quotes on the value
fmt.Fprintf(&sb, " (default %q)", f.DefValue)
} else {
fmt.Fprintf(&sb, " (default %v)", f.DefValue)
}
}
sb.WriteString("\n")
})
return sb.String()
}
// rpad adds padding to the right of a string.
func rpad(s string, padding int) string {
template := fmt.Sprintf("%%-%ds", padding)
return fmt.Sprintf(template, s)
}
func trimRightSpace(s string) string {
return strings.TrimRightFunc(s, unicode.IsSpace)
}
type PositionalArgs func(cmd *Command, args []string) error
// MinArgs returns an error if there is not at least N args.
func MinArgs(n int) PositionalArgs {
return func(cmd *Command, args []string) error {
if len(args) < n {
return fmt.Errorf("requires at least %d arg(s), only received %d", n, len(args))
}
return nil
}
}
// MaxArgs returns an error if there are more than N args.
func MaxArgs(n int) PositionalArgs {
return func(cmd *Command, args []string) error {
if len(args) > n {
return fmt.Errorf("accepts at most %d arg(s), received %d", n, len(args))
}
return nil
}
}
// ExactArgs returns an error if there are not exactly n args.
func ExactArgs(n int) PositionalArgs {
return func(cmd *Command, args []string) error {
if len(args) != n {
return fmt.Errorf("accepts %d arg(s), received %d", n, len(args))
}
return nil
}
}
// RangeArgs returns an error if the number of args is not within the expected range.
func RangeArgs(min int, max int) PositionalArgs {
return func(cmd *Command, args []string) error {
if len(args) < min || len(args) > max {
return fmt.Errorf("accepts between %d and %d arg(s), received %d", min, max, len(args))
}
return nil
}
}
func Execute(ctx context.Context, root *Command, args []string) error {
var (
showVersion bool
)
if root.Version != "" {
root.Flags().BoolVar(&showVersion, "v", false, "show version")
}
cmd, n := root.Find(args)
f := cmd.Flags()
if f != nil {
if err := f.Parse(args[n:]); err != nil {
if err == flag.ErrHelp {
t := template.Must(template.New("help").Funcs(templateFuncs).Parse(HelpTemplate))
return t.Execute(os.Stdout, &CommandHelp{cmd})
}
return err
}
}
if showVersion {
fmt.Println(root.Version)
return nil
}
if cmd.Args != nil {
if err := cmd.Args(cmd, f.Args()); err != nil {
return err
}
}
cmd.Run(ctx, f.Args())
return nil
}
// isZeroValue determines whether the string represents the zero
// value for a flag.
func isZeroValue(f *flag.Flag, value string) bool {
// Build a zero value of the flag's Value type, and see if the
// result of calling its String method equals the value passed in.
// This works unless the Value type is itself an interface type.
typ := reflect.TypeOf(f.Value)
var z reflect.Value
if typ.Kind() == reflect.Ptr {
z = reflect.New(typ.Elem())
} else {
z = reflect.Zero(typ)
}
return value == z.Interface().(flag.Value).String()
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment