command/command.go

package command

import (
	"errors"
	"fmt"
	"os"
	"reflect"
	"strconv"
	"strings"
	"text/template"
)

type Option struct {
	Name         string
	LongName     string
	ShortName    string
	Description  string
	OType        string
	DefaultValue interface{}
}

var UsageTemplate, _ = template.New("Usage").Parse(`Usage:
 {{if .ParentCommand }} {{.ParentCommand.Name}} {{end}} {{.Name}} [OPTIONS] [ARGS]
{{ if .Description }}
Description:
  {{.Description}}
{{ end }}
{{ if .Args }}
Options:
  {{range .Args}}
  -{{.ShortName}}|--{{.LongName}} {{.OType}} {{.Description}} default:{{.DefaultValue}}
  {{end}}
{{ end }}
{{ if .SubCommands }}
 Sub Commands:
  {{range .SubCommands}}
  {{.Name}} {{.Description}}
  {{end}}
{{end}}
`)

type Command struct {
	Name              string
	Description       string
	Args              []*Option
	SubCommands       []*Command
	ParentCommand     *Command
	subcommand        string
	globalOptions     map[string]interface{}
	subcommandOptions map[string]interface{}
	remainArgs        []string
}

func (c *Command) GetGlobalOptions() map[string]interface{} {
	return c.globalOptions
}
func (c *Command) GetSubCommandOptions() map[string]interface{} {
	return c.subcommandOptions
}
func (c *Command) GetSubCommand() string {
	return c.subcommand
}
func (c *Command) AddArg(name string, shortName string, description string, defaultValue interface{}) {
	oType := reflect.TypeOf(defaultValue).String()
	c.Args = append(c.Args, &Option{
		Name:         name,
		LongName:     ToKebabCase(name),
		ShortName:    shortName,
		Description:  description,
		OType:        oType,
		DefaultValue: defaultValue,
	})
}

func (c *Command) AddSubCommand(name string, description string) *Command {
	if c.ParentCommand != nil {
		panic("Sub commands can only be added to top level commands")
	}
	command := &Command{
		Name:          name,
		Description:   description,
		ParentCommand: c,
	}
	c.SubCommands = append(c.SubCommands, command)
	return command
}
func (c *Command) Usage() {
	UsageTemplate.Execute(os.Stdout, c)
}
func (c *Command) GetOption(name string) *Option {
	if len(name) > 1 {
		for _, arg := range c.Args {
			if arg.LongName == name {
				return arg
			}
		}
	}
	for _, arg := range c.Args {
		if arg.ShortName == name {
			return arg
		}
	}
	return nil
}
func (c *Command) parseError(errMsg string) {
	fmt.Println(errMsg)
	os.Exit(1)
}
func (c *Command) showHelpWithOption(args []string) {
	if len(args) > 0 && (args[0] == "-h" || args[0] == "--help") {
		c.Usage()
		os.Exit(0)
	}
}
func (c *Command) showSubCommandHelp(args []string) {
	if len(args) == 2 && args[0] == "help" {
		cmd := c.findSubcommand(args[1])
		if cmd == nil {
			c.parseError("Unknown subcommand " + args[1])
		}
		cmd.Usage()
		os.Exit(0)
	}
}
func (c *Command) Parse(args []string) ([]string, error) {
	c.showHelpWithOption(args)
	c.showSubCommandHelp(args)
	if len(args) == 0 {
		if len(c.Args) != 0 || len(c.SubCommands) != 0 || c.ParentCommand == nil {
			c.Usage()
		}
		if c.ParentCommand != nil {
			c.ParentCommand.subcommand = c.Name
		}
	}

	vargs := args
	for {
		if len(vargs) == 0 {
			break
		}
		var parsed bool
		var err error
		vargs, parsed, err = c.longOption(vargs)
		if err != nil {
			c.parseError(err.Error())
		}
		if parsed {
			continue
		}
		vargs, parsed, err = c.shortOption(vargs)
		if err != nil {
			c.parseError(err.Error())
		}
		if parsed {
			continue
		}
		if len(c.SubCommands) > 0 {
			cmd := c.findSubcommand(args[0])
			if cmd != nil {
				c.subcommand = cmd.Name
				vargs = vargs[1:]
				vargs, err = cmd.Parse(vargs)
				if err != nil {
					c.parseError(err.Error())
				}
			} else {
				c.parseError("Unknown subcommand " + args[0])
			}
		} else {
			if c.ParentCommand != nil {
				c.ParentCommand.remainArgs = append(c.ParentCommand.remainArgs, vargs[0])

			} else {
				c.remainArgs = append(c.remainArgs, vargs[0])
			}
			vargs = vargs[1:]

		}
	}
	return vargs, nil
}
func (c *Command) findSubcommand(cmd string) *Command {
	for _, subCmd := range c.SubCommands {
		if subCmd.Name == cmd {
			return subCmd
		}
	}
	return nil
}

func (c *Command) longOption(args []string) ([]string, bool, error) {
	if !strings.HasPrefix(args[0], "--") {
		return args, false, nil
	}
	optName := strings.TrimPrefix(args[0], "--")
	if len(optName) < 2 {
		return args, false, errors.New("Invalid option name")
	}
	opt := c.GetOption(optName)
	if opt != nil {
		rargs, _, err := c.getOptValue(args, opt, false)
		if err == nil {
			return rargs, true, nil
		} else {
			return args, false, err
		}
	}
	if c.ParentCommand != nil {
		return c.ParentCommand.longOption(args)
	}
	return args, false, errors.New("Unknown option " + args[0])
}
func (c *Command) shortOption(args []string) ([]string, bool, error) {
	if !strings.HasPrefix(args[0], "-") {
		return args, false, nil
	}
	optName := strings.TrimPrefix(args[0], "-")
	last := len(optName) - 1
	for i, s := range optName {
		opt := c.GetOption(string(s))
		if opt != nil {
			if i != last {
				_, _, err := c.getOptValue(args[i:], opt, true)
				if err == nil {
					return args, true, nil
				} else {
					continue
				}
			} else {
				rargs, _, err := c.getOptValue(args[i:], opt, false)
				if err == nil {
					return rargs, true, nil
				} else {
					return args, false, err
				}
			}

		} else {
			if c.ParentCommand != nil {
				opt = c.ParentCommand.GetOption(string(s))
				if opt != nil {
					rargs, _, err := c.ParentCommand.getOptValue(args[i:], opt, i != last)
					if err == nil {
						return rargs, true, nil
					} else {
						return args, false, err
					}
				}
			} else {
				return args, false, errors.New("Unknown option " + args[0])
			}
		}
	}
	return args, false, errors.New("Unknown option " + args[0])
}

func (c *Command) getOptValue(args []string, opt *Option, isFixBool bool) ([]string, interface{}, error) {

	paramMap := c.globalOptions
	if c.ParentCommand != nil {
		paramMap = c.ParentCommand.subcommandOptions
	}
	if isFixBool {
		if opt.OType != "bool" {
			return args, nil, errors.New("Invalid type for boolean option")
		}
		paramMap[opt.Name] = true
		return args, true, nil
	}

	if opt.OType == "bool" {
		nextStr := strings.ToLower(args[1])
		if nextStr == "f" || nextStr == "false" {
			paramMap[opt.Name] = false
			return args[2:], true, nil
		} else if nextStr == "t" || nextStr == "true" {
			paramMap[opt.Name] = true
			return args[2:], true, nil
		} else {
			paramMap[opt.Name] = true
			return args[1:], true, nil
		}
	} else {
		valStr := args[1]
		switch opt.OType {
		case "int":
			val, err := strconv.Atoi(valStr)
			if err != nil {
				return args, nil, errors.New("Invalid value for option " + opt.Name)
			} else {
				paramMap[opt.Name] = val
				return args[2:], true, nil
			}
		case "string":
			paramMap[opt.Name] = valStr
			return args[2:], true, nil
		case "float64":
			val, err := strconv.ParseFloat(valStr, 64)
			if err != nil {
				return args, nil, errors.New("Invalid value for option " + opt.Name)
			} else {
				paramMap[opt.Name] = val
				return args[2:], true, nil
			}
		case "default":
			return args, nil, errors.New("unsupported type")
		}
	}
	return args, nil, errors.New("unsupported type")
}

func NewCommand(name string, desc string) *Command {
	return &Command{
		Name:              name,
		Description:       desc,
		globalOptions:     map[string]interface{}{},
		subcommandOptions: map[string]interface{}{},
		remainArgs:        []string{},
	}
}