feat: add completion install api

Author: ethanndickson

command.go

@@ -296,10 +296,16 @@ func copyFlagSetWithout(fs *pflag.FlagSet, without string) *pflag.FlagSet {
 }
 
 func (inv *Invocation) CurWords() (prev string, cur string) {
-	if len(inv.Args) == 1 {
+	switch len(inv.Args) {
+	// All the shells we support will supply at least one argument (empty string),
+	// but we don't want to panic.
+	case 0:
+		cur = ""
+		prev = ""
+	case 1:
 		cur = inv.Args[0]
 		prev = ""
-	} else {
+	default:
 		cur = inv.Args[len(inv.Args)-1]
 		prev = inv.Args[len(inv.Args)-2]
 	}
@@ -645,9 +651,13 @@ func (inv *Invocation) completeFlag(word string) []string {
 	if opt.CompletionHandler != nil {
 		return opt.CompletionHandler(inv)
 	}
-	val, ok := opt.Value.(*Enum)
+	enum, ok := opt.Value.(*Enum)
+	if ok {
+		return enum.Choices
+	}
+	enumArr, ok := opt.Value.(*EnumArray)
 	if ok {
-		return val.Choices
+		return enumArr.Choices
 	}
 	return nil
 }

command_test.go

@@ -34,14 +34,15 @@ func fakeIO(i *serpent.Invocation) *ioBufs {
 func sampleCommand(t *testing.T) *serpent.Command {
 	t.Helper()
 	var (
-		verbose bool
-		lower   bool
-		prefix  string
-		reqBool bool
-		reqStr  string
-		reqArr  []string
-		fileArr []string
-		enumStr string
+		verbose    bool
+		lower      bool
+		prefix     string
+		reqBool    bool
+		reqStr     string
+		reqArr     []string
+		reqEnumArr []string
+		fileArr    []string
+		enumStr    string
 	)
 	enumChoices := []string{"foo", "bar", "qux"}
 	return &serpent.Command{
@@ -94,6 +95,11 @@ func sampleCommand(t *testing.T) *serpent.Command {
 						FlagShorthand: "a",
 						Value:         serpent.StringArrayOf(&reqArr),
 					},
+					serpent.Option{
+						Name:  "req-enum-array",
+						Flag:  "req-enum-array",
+						Value: serpent.EnumArrayOf(&reqEnumArr, enumChoices...),
+					},
 				},
 				HelpHandler: func(i *serpent.Invocation) error {
 					_, _ = i.Stdout.Write([]byte("help text.png"))

completion.go

@@ -1,5 +1,7 @@
 package serpent
 
+import "strings"
+
 // CompletionModeEnv is a special environment variable that is
 // set when the command is being run in completion mode.
 const CompletionModeEnv = "COMPLETION_MODE"
@@ -18,7 +20,9 @@ func DefaultCompletionHandler(inv *Invocation) []string {
 		allResps = append(allResps, cmd.Name())
 	}
 	for _, opt := range inv.Command.Options {
-		if opt.ValueSource == ValueSourceNone || opt.ValueSource == ValueSourceDefault || opt.Value.Type() == "string-array" {
+		if opt.ValueSource == ValueSourceNone ||
+			opt.ValueSource == ValueSourceDefault ||
+			strings.Contains(opt.Value.Type(), "array") {
 			allResps = append(allResps, "--"+opt.Flag)
 		}
 	}

completion/all.go

@@ -6,74 +6,81 @@ import (
 	"os"
 	"os/user"
 	"path/filepath"
+	"runtime"
 	"strings"
 	"text/template"
 
 	"github.com/coder/serpent"
 )
 
+type Shell interface {
+	Name() string
+	InstallPath() (string, error)
+	UsesOwnFile() bool
+	WriteCompletion(io.Writer) error
+}
+
 const (
-	BashShell  string = "bash"
-	FishShell  string = "fish"
-	ZShell     string = "zsh"
-	Powershell string = "powershell"
+	ShellBash       string = "bash"
+	ShellFish       string = "fish"
+	ShellZsh        string = "zsh"
+	ShellPowershell string = "powershell"
 )
 
-var shellCompletionByName = map[string]func(io.Writer, string) error{
-	BashShell:  generateCompletion(bashCompletionTemplate),
-	FishShell:  generateCompletion(fishCompletionTemplate),
-	ZShell:     generateCompletion(zshCompletionTemplate),
-	Powershell: generateCompletion(pshCompletionTemplate),
+func ShellByName(shell, programName string) (Shell, error) {
+	switch shell {
+	case ShellBash:
+		return Bash(runtime.GOOS, programName), nil
+	case ShellFish:
+		return Fish(runtime.GOOS, programName), nil
+	case ShellZsh:
+		return Zsh(runtime.GOOS, programName), nil
+	case ShellPowershell:
+		return Powershell(runtime.GOOS, programName), nil
+	default:
+		return nil, fmt.Errorf("unsupported shell %q", shell)
+	}
 }
 
 func ShellOptions(choice *string) *serpent.Enum {
-	return serpent.EnumOf(choice, BashShell, FishShell, ZShell, Powershell)
-}
-
-func WriteCompletion(writer io.Writer, shell string, cmdName string) error {
-	fn, ok := shellCompletionByName[shell]
-	if !ok {
-		return fmt.Errorf("unknown shell %q", shell)
-	}
-	fn(writer, cmdName)
-	return nil
+	return serpent.EnumOf(choice, ShellBash, ShellFish, ShellZsh, ShellPowershell)
 }
 
-func DetectUserShell() (string, error) {
+func DetectUserShell(programName string) (Shell, error) {
 	// Attempt to get the SHELL environment variable first
 	if shell := os.Getenv("SHELL"); shell != "" {
-		return filepath.Base(shell), nil
+		return ShellByName(filepath.Base(shell), "")
 	}
 
 	// Fallback: Look up the current user and parse /etc/passwd
 	currentUser, err := user.Current()
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
 	// Open and parse /etc/passwd
 	passwdFile, err := os.ReadFile("/etc/passwd")
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
 	lines := strings.Split(string(passwdFile), "\n")
 	for _, line := range lines {
 		if strings.HasPrefix(line, currentUser.Username+":") {
 			parts := strings.Split(line, ":")
 			if len(parts) > 6 {
-				return filepath.Base(parts[6]), nil // The shell is typically the 7th field
+				return ShellByName(filepath.Base(parts[6]), programName) // The shell is typically the 7th field
 			}
 		}
 	}
 
-	return "", fmt.Errorf("default shell not found")
+	return nil, fmt.Errorf("default shell not found")
 }
 
 func generateCompletion(
 	scriptTemplate string,
 ) func(io.Writer, string) error {
-	return func(w io.Writer, rootCmdName string) error {
+	return func(w io.Writer, programName string) error {
 		tmpl, err := template.New("script").Parse(scriptTemplate)
 		if err != nil {
 			return fmt.Errorf("parse template: %w", err)
@@ -82,7 +89,7 @@ func generateCompletion(
 		err = tmpl.Execute(
 			w,
 			map[string]string{
-				"Name": rootCmdName,
+				"Name": programName,
 			},
 		)
 		if err != nil {
@@ -92,3 +99,35 @@ func generateCompletion(
 		return nil
 	}
 }
+
+func InstallShellCompletion(shell Shell) error {
+	path, err := shell.InstallPath()
+	if err != nil {
+		return fmt.Errorf("get install path: %w", err)
+	}
+
+	err = os.MkdirAll(filepath.Dir(path), 0o755)
+	if err != nil {
+		return fmt.Errorf("create directories: %w", err)
+	}
+
+	if shell.UsesOwnFile() {
+		err := os.WriteFile(path, nil, 0o644)
+		if err != nil {
+			return fmt.Errorf("create file: %w", err)
+		}
+	}
+
+	f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
+	if err != nil {
+		return fmt.Errorf("open file for appending: %w", err)
+	}
+	defer f.Close()
+
+	err = shell.WriteCompletion(f)
+	if err != nil {
+		return fmt.Errorf("write completion script: %w", err)
+	}
+
+	return nil
+}

completion/bash.go

@@ -1,6 +1,53 @@
 package completion
 
+import (
+	"io"
+	"path/filepath"
+
+	home "github.com/mitchellh/go-homedir"
+)
+
+type bash struct {
+	goos        string
+	programName string
+}
+
+var _ Shell = &bash{}
+
+func Bash(goos string, programName string) Shell {
+	return &bash{goos: goos, programName: programName}
+}
+
+// Name implements Shell.
+func (b *bash) Name() string {
+	return "bash"
+}
+
+// UsesOwnFile implements Shell.
+func (b *bash) UsesOwnFile() bool {
+	return false
+}
+
+// InstallPath implements Shell.
+func (b *bash) InstallPath() (string, error) {
+	homeDir, err := home.Dir()
+	if err != nil {
+		return "", err
+	}
+	if b.goos == "darwin" {
+		return filepath.Join(homeDir, ".bash_profile"), nil
+	}
+	return filepath.Join(homeDir, ".bashrc"), nil
+}
+
+// WriteCompletion implements Shell.
+func (b *bash) WriteCompletion(w io.Writer) error {
+	return generateCompletion(bashCompletionTemplate)(w, b.programName)
+}
+
 const bashCompletionTemplate = `
+
+# === BEGIN {{.Name}} COMPLETION ===
 _generate_{{.Name}}_completions() {
     # Capture the line excluding the command, and everything after the current word
     local args=("${COMP_WORDS[@]:1:COMP_CWORD}")
@@ -16,7 +63,8 @@ _generate_{{.Name}}_completions() {
         COMPREPLY=()
     fi
 }
-
 # Setup Bash to use the function for completions for '{{.Name}}'
 complete -F _generate_{{.Name}}_completions {{.Name}}
+# === END {{.Name}} COMPLETION ===
+
 `

completion/fish.go

@@ -1,5 +1,47 @@
 package completion
 
+import (
+	"io"
+	"path/filepath"
+
+	home "github.com/mitchellh/go-homedir"
+)
+
+type fish struct {
+	goos        string
+	programName string
+}
+
+var _ Shell = &fish{}
+
+func Fish(goos string, programName string) Shell {
+	return &fish{goos: goos, programName: programName}
+}
+
+// UsesOwnFile implements Shell.
+func (f *fish) UsesOwnFile() bool {
+	return true
+}
+
+// Name implements Shell.
+func (f *fish) Name() string {
+	return "fish"
+}
+
+// InstallPath implements Shell.
+func (f *fish) InstallPath() (string, error) {
+	homeDir, err := home.Dir()
+	if err != nil {
+		return "", err
+	}
+	return filepath.Join(homeDir, ".config/fish/completions/", f.programName+".fish"), nil
+}
+
+// WriteCompletion implements Shell.
+func (f *fish) WriteCompletion(w io.Writer) error {
+	return generateCompletion(fishCompletionTemplate)(w, f.programName)
+}
+
 const fishCompletionTemplate = `
 function _{{.Name}}_completions
 	# Capture the full command line as an array