pmt/cmd/command.go

package cmd

import (
	"errors"
	"fmt"
	"os"
	"strings"
)

type DyadicValue struct {
	Value   any
	Another any
}

func (d *DyadicValue) GetValue() any {
	if d.Value != nil {
		return d.Value
	}
	return d.Another
}

type CommandAction func(*Command, map[string]any, ...any) error

type Command struct {
	name            string
	usage           *DyadicValue
	version         *Option
	help            *Option
	description     string
	longDescription string
	arguments       []*Argument
	options         []*Option
	types           map[string]TypeParser
	globalTypes     map[string]TypeParser
	commands        []*Command
	program         *Command
	action          CommandAction
}

func New(name string, desc ...string) *Command {
	cmd := &Command{
		types: make(map[string]TypeParser),
	}
	cmd.Version("1.0.0").Help("")
	switch len(desc) {
	case 0:
		return cmd.Name(name)
	case 1:
		return cmd.Name(name).Description(desc[0])
	default:
		return cmd.Name(name).Description(desc[0], desc[1:]...)
	}
}

func (c *Command) Name(name string) *Command {
	if len(name) == 0 {
		panic(errors.New("empty command name"))
	}

	// 仅仅支持二级，所以只需要检查 Command.program.commands 中是否重复
	if c.program != nil {
		for _, command := range c.program.commands {
			if command.name == name {
				panic(errors.New("duplicated command name"))
			}
		}
	}

	c.name = name

	return c
}

func (c *Command) GetName() string {
	return c.name
}

func (c *Command) Type(name string, parser func(...string) any) *Command {
	c.types[name] = parser
	return c
}

func (c *Command) GlobalType(name string, parser func(...string) any) *Command {
	if c.program != nil {
		c.program.GlobalType(name, parser)
	} else {
		c.globalTypes[name] = parser
	}
	return c
}

func (c *Command) Command(name string, desc ...string) *Command {
	program := c.program
	if program == nil {
		program = c
	}

	cmd := New(name, desc...)
	cmd.program = program

	program.commands = append(program.commands, cmd)

	return cmd
}

func (c *Command) Version(ver string) *Command {
	if c.version == nil {
		c.VersionOption("-v, --version", "Show version number")
	}
	if len(ver) > 0 {
		c.version.RuntimeValue = []string{ver}
	} else {
		c.version.RuntimeValue = make([]string, 0)
	}
	return c
}

func (c *Command) VersionOption(expr string, desc ...string) *Command {
	ver, err := ensureOption(c, expr, desc, "version")
	if errors.Is(err, ErrOptionDefinition) {
		panic(ErrVersionDefinition)
	} else if err != nil {
		panic(err)
	}
	if c.help != nil {
		err = hasConflicts(c.help, ver)
		if err != nil {
			panic(errors.New("version option conflicts with help option, " + err.Error()))
		}
	}
	c.version = ver
	return c
}

func (c *Command) GetVersion(opts ...*ShowOptions) string {
	if c.program != nil {
		return c.program.GetVersion()
	}
	if c.version == nil {
		c.Version("1.0.0")
	}
	return versionToString(c, resolveShowOptions(opts))
}

func (c *Command) ShowVersion(opts ...*ShowOptions) error {
	return showMessage(resolveShowOptions(opts), c.GetVersion)
}

func (c *Command) Usage(usage string) *Command {
	if c.usage == nil {
		c.usage = new(DyadicValue)
	}
	c.usage.Value = usage
	return c
}

func (c *Command) GetUsage(opts ...*ShowOptions) string {
	if c.usage != nil {
		usage := c.usage.GetValue()
		if usage != nil {
			return usage.(string)
		}
	}
	s := usageToString(c, resolveShowOptions(opts))
	c.usage = new(DyadicValue)
	c.usage.Another = s
	return s
}

func (c *Command) ShowUsage(opts ...*ShowOptions) error {
	return showMessage(resolveShowOptions(opts), c.GetUsage)
}

func (c *Command) Description(desc string, longDesc ...string) *Command {
	c.description = desc
	if len(longDesc) > 0 {
		c.longDescription = strings.Join(longDesc, "\n")
	}
	return c
}

func (c *Command) GetDescription(opts ...*ShowOptions) string {
	o := resolveShowOptions(opts)
	var s string
	if o.Long {
		s = c.longDescription
	} else {
		s = c.description
	}
	if len(s) == 0 {
		return ""
	}
	if o.Titles {
		return "DESCRIPTION:\n    " + s
	}
	return s
}

func (c *Command) ShowDescription(opts ...*ShowOptions) error {
	return showMessage(resolveShowOptions(opts), c.GetDescription)
}

func (c *Command) Help(help string) *Command {
	if c.help == nil {
		c.HelpOption("-h, --help", "Show help information")
	}
	if len(help) > 0 {
		c.help.RuntimeValue = []string{help}
	} else {
		c.help.RuntimeValue = make([]string, 0)
	}
	return c
}

func (c *Command) HelpOption(expr string, desc ...string) *Command {
	help, err := ensureOption(c, expr, desc, "help")
	if errors.Is(err, ErrOptionDefinition) {
		panic(ErrVersionDefinition)
	} else if err != nil {
		panic(err)
	}
	if c.version != nil {
		err = hasConflicts(help, c.version)
		if err != nil {
			panic(errors.New("help option conflicts with version option, " + err.Error()))
		}
	}
	c.help = help
	return c
}

func (c *Command) GetHelp(opts ...*ShowOptions) string {
	if c.help != nil {
		if len(c.help.RuntimeValue) > 0 {
			return c.help.RuntimeValue[0]
		}
		if c.help.DefaultValue != nil {
			return c.help.DefaultValue.(string)
		}
	}
	s := helpToString(c, resolveShowOptions(opts))
	c.help.DefaultValue = s
	return s
}

func (c *Command) ShowHelp(opts ...*ShowOptions) error {
	return showMessage(resolveShowOptions(opts), c.GetHelp)
}

func (c *Command) Argument(expr string, desc ...string) *Command {
	a, err := createArgument(expr)
	if err != nil {
		panic(err)
	}
	switch len(desc) {
	case 0:
		break
	case 1:
		a.Description = desc[0]
	default:
		a.Description = desc[0]
		a.LongDescription = desc[1]
	}
	var lastArg *Argument
	for _, a2 := range c.arguments {
		lastArg = a2
		if a2.Name == a.Name {
			panic(errors.New("duplicated arguments definition"))
		}
	}
	if lastArg != nil {
		if lastArg.Variadic {
			panic(errors.New("after variadic argument"))
		}
		if a.Required && !lastArg.Required {
			panic(errors.New("after optional argument"))
		}
	}
	c.arguments = append(c.arguments, a)
	return c
}

func (c *Command) Option(expr string, desc ...string) *Command {
	o, err := ensureOption(c, expr, desc, "")
	if err != nil {
		panic(err)
	}
	if c.version != nil {
		err = hasConflicts(c.version, o)
		if err != nil {
			panic(err)
		}
	}
	if c.help != nil {
		err = hasConflicts(c.help, o)
		if err != nil {
			panic(err)
		}
	}
	c.options = append(c.options, o)
	return c
}

func (c *Command) Action(action func(*Command, map[string]any, ...any) error) *Command {
	c.action = action
	return c
}

func (c *Command) Run(args ...[]string) error {
	var argv []string
	for _, arg := range args {
		argv = arg
		break
	}
	if argv == nil {
		argv = os.Args[1:]
	}

	if c.program != nil {
		return c.program.Run(argv)
	}

	command := c
	cmdIdx := -1
	for i, arg := range argv {
		if !strings.HasPrefix(argv[0], "-") {
			for _, cmd := range c.commands {
				if cmd.name == arg {
					command = cmd
					cmdIdx = i
					break
				}
			}
			break
		}
	}
	if cmdIdx == 0 {
		argv = argv[1:]
	} else if cmdIdx == len(argv)-1 {
		argv = argv[:cmdIdx]
	} else if cmdIdx > -1 {
		argv = append(argv[:cmdIdx], argv[cmdIdx+1:]...)
	}

	help, version, options, arguments := parseArgs(argv, command)
	opts := resolveShowOptions([]*ShowOptions{})

	if help {
		return command.ShowHelp(opts)
	} else if version {
		return command.ShowVersion(&ShowOptions{
			Colors:    opts.Colors,
			Types:     opts.Types,
			Long:      opts.Long,
			Titles:    false,
			LineWidth: opts.LineWidth,
			Writer:    opts.Writer,
		})
	} else if command.action == nil {
		fmt.Println("没有操作，则显示帮助文档")
		// 没有操作，则显示帮助文档
		return command.ShowHelp()
	} else {
		return command.action(command, options, arguments...)
	}
}

func hasConflicts(o1, o2 *Option) error {
	if len(o1.Short) > 0 && o1.Short == o2.Short {
		return errors.New("duplicated short flag of option definition")
	}
	if o1.Long == o2.Long {
		return errors.New("duplicated long flag of option definition")
	}
	if o1.Name == o2.Name {
		return errors.New("duplicated arguments definition")
	}
	return nil
}