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 in 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
@@ -193,15 +199,24 @@ 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
+	// 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)
@@ -282,6 +297,17 @@ func copyFlagSetWithout(fs *pflag.FlagSet, without string) *pflag.FlagSet {
 	return fs2
 }
 
+func (inv *Invocation) curWords() (prev string, cur string) {
+	if len(inv.Args) == 1 {
+		cur = inv.Args[0]
+		prev = ""
+	} else {
+		cur = inv.Args[len(inv.Args)-1]
+		prev = inv.Args[len(inv.Args)-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 +404,19 @@ func (inv *Invocation) run(state *runState) error {
 		}
 	}
 
+	// 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() {
+		for _, e := range inv.doCompletions() {
+			fmt.Fprintln(inv.Stdout, e)
+		}
+		return nil
+	}
+
+	ignoreFlagParseErrors := inv.Command.RawArgs
+
 	// 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 +433,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, ", "))
 	}
 
@@ -553,6 +590,78 @@ func (inv *Invocation) with(fn func(*Invocation)) *Invocation {
 	return &i2
 }
 
+func (inv *Invocation) doCompletions() []string {
+	prev, cur := inv.curWords()
+	inv.CurWord = cur
+	// If the current word is a flag set using `=`, use it's handler
+	if strings.HasPrefix(cur, "--") && strings.Contains(cur, "=") {
+		if out := inv.equalsFlagCompletions(cur); out != nil {
+			return out
+		}
+	}
+	// 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 out := inv.flagCompletions(prev); out != nil {
+			return out
+		}
+	}
+	// If the current word is the command, auto-complete it so the shell moves the cursor
+	if inv.Command.Name() == inv.CurWord {
+		return []string{inv.Command.Name()}
+	}
+	var completions []string
+
+	if inv.Command.CompletionHandler != nil {
+		completions = append(completions, inv.Command.CompletionHandler(inv)...)
+	}
+
+	completions = append(completions, DefaultCompletionHandler(inv)...)
+
+	return completions
+}
+
+func (inv *Invocation) flagCompletions(word string) []string {
+	return inv.doFlagCompletions("", word)
+}
+
+func (inv *Invocation) equalsFlagCompletions(word string) []string {
+	words := strings.Split(word, "=")
+	word = words[0]
+	if len(words) > 1 {
+		inv.CurWord = words[1]
+	} else {
+		inv.CurWord = ""
+	}
+	prefix := word + "="
+	return inv.doFlagCompletions(prefix, word)
+}
+
+func (inv *Invocation) doFlagCompletions(prefix, word string) []string {
+	opt := inv.Command.Options.ByFlag(word[2:])
+	if opt == nil {
+		return nil
+	}
+	if opt.CompletionHandler != nil {
+		completions := opt.CompletionHandler(inv)
+		out := make([]string, 0, len(completions))
+		for _, completion := range completions {
+			out = append(out, fmt.Sprintf("%s%s", prefix, completion))
+		}
+		return out
+	}
+	val, ok := opt.Value.(*Enum)
+	if ok {
+		completions := val.Choices
+		out := make([]string, 0, len(completions))
+		for _, choice := range completions {
+			out = append(out, fmt.Sprintf("%s%s", prefix, choice))
+		}
+		return out
+	}
+	return nil
+}
+
 // MiddlewareFunc returns the next handler in the chain,
 // or nil if there are no more.
 type MiddlewareFunc func(next HandlerFunc) HandlerFunc
@@ -637,3 +746,5 @@ 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