Add auto-completion

command.go

@@ -59,6 +59,12 @@ type Command struct {
 	Middleware  MiddlewareFunc
 	Handler     HandlerFunc
 	HelpHandler HandlerFunc
+	// CompletionHandler is called when the command is run is completion
+	// mode. If nil, only the default completion handler is used.
+	//
+	// Flag and option parsing is best-effort in this mode, so even if an Option
+	// is "required" it may not be set.
+	CompletionHandler CompletionHandlerFunc
 }
 
 // AddSubcommands adds the given subcommands, setting their
@@ -181,6 +187,7 @@ func (c *Command) Invoke(args ...string) *Invocation {
 	return &Invocation{
 		Command: c,
 		Args:    args,
+		AllArgs: args,
 		Stdout:  io.Discard,
 		Stderr:  io.Discard,
 		Stdin:   strings.NewReader(""),
@@ -193,15 +200,27 @@ type Invocation struct {
 	ctx         context.Context
 	Command     *Command
 	parsedFlags *pflag.FlagSet
-	Args        []string
+
+	// Args is reduced into the remaining arguments after parsing flags
+	// during Run.
+	Args []string
+	// AllArgs is the original arguments passed to the command, including flags.
+	// When invoked `WithOS`, this includes argv[0], otherwise it is the same as Args.
+	AllArgs []string
+	// CurWord is the word the terminal cursor is currently in
+	CurWord string
+
 	// Environ is a list of environment variables. Use EnvsWithPrefix to parse
 	// os.Environ.
 	Environ Environ
 	Stdout  io.Writer
 	Stderr  io.Writer
 	Stdin   io.Reader
-	Logger  slog.Logger
-	Net     Net
+
+	// Deprecated
+	Logger slog.Logger
+	// Deprecated
+	Net Net
 
 	// testing
 	signalNotifyContext func(parent context.Context, signals ...os.Signal) (ctx context.Context, stop context.CancelFunc)
@@ -214,6 +233,7 @@ func (inv *Invocation) WithOS() *Invocation {
 		i.Stdout = os.Stdout
 		i.Stderr = os.Stderr
 		i.Stdin = os.Stdin
+		i.AllArgs = os.Args
 		i.Args = os.Args[1:]
 		i.Environ = ParseEnviron(os.Environ(), "")
 		i.Net = osNet{}
@@ -282,6 +302,17 @@ func copyFlagSetWithout(fs *pflag.FlagSet, without string) *pflag.FlagSet {
 	return fs2
 }
 
+func (inv *Invocation) GetCurWords() (prev string, cur string) {
+	if len(inv.AllArgs) == 1 {
+		cur = inv.AllArgs[0]
+		prev = ""
+	} else {
+		cur = inv.AllArgs[len(inv.AllArgs)-1]
+		prev = inv.AllArgs[len(inv.AllArgs)-2]
+	}
+	return
+}
+
 // run recursively executes the command and its children.
 // allArgs is wired through the stack so that global flags can be accepted
 // anywhere in the command invocation.
@@ -378,8 +409,10 @@ func (inv *Invocation) run(state *runState) error {
 		}
 	}
 
+	ignoreFlagParseErrors := inv.Command.RawArgs || inv.IsCompletionMode()
+
 	// Flag parse errors are irrelevant for raw args commands.
-	if !inv.Command.RawArgs && state.flagParseErr != nil && !errors.Is(state.flagParseErr, pflag.ErrHelp) {
+	if !ignoreFlagParseErrors && state.flagParseErr != nil && !errors.Is(state.flagParseErr, pflag.ErrHelp) {
 		return xerrors.Errorf(
 			"parsing flags (%v) for %q: %w",
 			state.allArgs,
@@ -396,7 +429,7 @@ func (inv *Invocation) run(state *runState) error {
 		}
 	}
 	// Don't error for missing flags if `--help` was supplied.
-	if len(missing) > 0 && !errors.Is(state.flagParseErr, pflag.ErrHelp) {
+	if len(missing) > 0 && !inv.IsCompletionMode() && !errors.Is(state.flagParseErr, pflag.ErrHelp) {
 		return xerrors.Errorf("Missing values for the required flags: %s", strings.Join(missing, ", "))
 	}
 
@@ -431,6 +464,39 @@ func (inv *Invocation) run(state *runState) error {
 	defer cancel()
 	inv = inv.WithContext(ctx)
 
+	// Outputted completions are not filtered based on the word under the cursor, as every shell we support does this already.
+	// We only look at the current word to figure out handler to run, or what directory to inspect.
+	if inv.IsCompletionMode() {
+		prev, cur := inv.GetCurWords()
+		inv.CurWord = cur
+		// If the current word is a flag set using `=`, use it's handler
+		if strings.HasPrefix(cur, "--") && strings.Contains(cur, "=") {
+			if inv.equalsFlagHandler(cur) {
+				return nil
+			}
+		}
+		// If the previous word is a flag, then we're writing it's value
+		// and we should check it's handler
+		if strings.HasPrefix(prev, "--") {
+			if inv.flagHandler(prev) {
+				return nil
+			}
+		}
+		if inv.Command.Name() == inv.CurWord {
+			fmt.Fprintf(inv.Stdout, "%s\n", inv.Command.Name())
+			return nil
+		}
+		if inv.Command.CompletionHandler != nil {
+			for _, e := range inv.Command.CompletionHandler(inv) {
+				fmt.Fprintf(inv.Stdout, "%s\n", e)
+			}
+		}
+		for _, e := range DefaultCompletionHandler(inv) {
+			fmt.Fprintf(inv.Stdout, "%s\n", e)
+		}
+		return nil
+	}
+
 	if inv.Command.Handler == nil || errors.Is(state.flagParseErr, pflag.ErrHelp) {
 		if inv.Command.HelpHandler == nil {
 			return defaultHelpFn()(inv)
@@ -553,6 +619,44 @@ func (inv *Invocation) with(fn func(*Invocation)) *Invocation {
 	return &i2
 }
 
+func (inv *Invocation) flagHandler(word string) bool {
+	return inv.doFlagCompletion("", word)
+}
+
+func (inv *Invocation) equalsFlagHandler(word string) bool {
+	words := strings.Split(word, "=")
+	word = words[0]
+	if len(words) > 1 {
+		inv.CurWord = words[1]
+	} else {
+		inv.CurWord = ""
+	}
+	prefix := word + "="
+	return inv.doFlagCompletion(prefix, word)
+}
+
+func (inv *Invocation) doFlagCompletion(prefix, word string) bool {
+	opt := inv.Command.Options.ByFlag(word[2:])
+	if opt == nil {
+		return false
+	}
+	if opt.CompletionHandler != nil {
+		completions := opt.CompletionHandler(inv)
+		for _, completion := range completions {
+			fmt.Fprintf(inv.Stdout, "%s%s\n", prefix, completion)
+		}
+		return true
+	}
+	val, ok := opt.Value.(*Enum)
+	if ok {
+		for _, choice := range val.Choices {
+			fmt.Fprintf(inv.Stdout, "%s%s\n", prefix, choice)
+		}
+		return true
+	}
+	return false
+}
+
 // MiddlewareFunc returns the next handler in the chain,
 // or nil if there are no more.
 type MiddlewareFunc func(next HandlerFunc) HandlerFunc
@@ -637,3 +741,7 @@ func RequireRangeArgs(start, end int) MiddlewareFunc {
 
 // HandlerFunc handles an Invocation of a command.
 type HandlerFunc func(i *Invocation) error
+
+type CompletionHandlerFunc func(i *Invocation) []string
+
+var NopHandler HandlerFunc = func(i *Invocation) error { return nil }