Skip to content

Commit

Permalink
feat: add completion install api
Browse files Browse the repository at this point in the history
  • Loading branch information
ethanndickson committed Aug 9, 2024
1 parent 91966a2 commit 460dd41
Show file tree
Hide file tree
Showing 14 changed files with 335 additions and 54 deletions.
8 changes: 6 additions & 2 deletions command.go
Original file line number Diff line number Diff line change
Expand Up @@ -645,9 +645,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 val.Choices
return enum.Choices
}
enumArr, ok := opt.Value.(*EnumArray)
if ok {
return enumArr.Choices
}
return nil
}
Expand Down
22 changes: 14 additions & 8 deletions command_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down Expand Up @@ -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"))
Expand Down
6 changes: 5 additions & 1 deletion completion.go
Original file line number Diff line number Diff line change
@@ -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"
Expand All @@ -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)
}
}
Expand Down
74 changes: 47 additions & 27 deletions completion/all.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,74 +6,80 @@ import (
"os"
"os/user"
"path/filepath"
"runtime"
"strings"
"text/template"

"github.com/coder/serpent"
)

type Shell interface {
Name() string
InstallPath() (string, error)
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 {
switch shell {
case ShellBash:
return Bash(runtime.GOOS, programName)
case ShellFish:
return Fish(runtime.GOOS, programName)
case ShellZsh:
return Zsh(runtime.GOOS, programName)
case ShellPowershell:
return Powershell(runtime.GOOS, programName)
default:
return nil
}
}

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), ""), nil
}

// 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), nil // 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)
Expand All @@ -82,7 +88,7 @@ func generateCompletion(
err = tmpl.Execute(
w,
map[string]string{
"Name": rootCmdName,
"Name": programName,
},
)
if err != nil {
Expand All @@ -92,3 +98,17 @@ func generateCompletion(
return nil
}
}

func InstallShellCompletion(shell Shell) error {
path, err := shell.InstallPath()
if err != nil {
return fmt.Errorf("get install path: %w", err)
}
f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
if err != nil {
return fmt.Errorf("create and append to file: %w", err)
}
defer f.Close()

return shell.WriteCompletion(f)
}
43 changes: 42 additions & 1 deletion completion/bash.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,47 @@
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"
}

// 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}")
Expand All @@ -16,7 +57,7 @@ _generate_{{.Name}}_completions() {
COMPREPLY=()
fi
}
# Setup Bash to use the function for completions for '{{.Name}}'
complete -F _generate_{{.Name}}_completions {{.Name}}
# === END {{.Name}} COMPLETION ===
`
37 changes: 37 additions & 0 deletions completion/fish.go
Original file line number Diff line number Diff line change
@@ -1,5 +1,42 @@
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}
}

// 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
Expand Down
48 changes: 46 additions & 2 deletions completion/powershell.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,51 @@
package completion

const pshCompletionTemplate = `
import (
"io"
"os/exec"
)

type powershell struct {
goos string
programName string
}

// Name implements Shell.
func (p *powershell) Name() string {
return "powershell"
}

func Powershell(goos string, programName string) Shell {
return &powershell{goos: goos, programName: programName}
}

// InstallPath implements Shell.
func (p *powershell) InstallPath() (string, error) {
var (
path []byte
err error
)
cmd := "$PROFILE.CurrentUserAllHosts"
if p.goos == "windows" {
path, err = exec.Command("powershell", cmd).CombinedOutput()
} else {
path, err = exec.Command("pwsh", "-Command", cmd).CombinedOutput()
}
if err != nil {
return "", err
}
return string(path), nil
}

// WriteCompletion implements Shell.
func (p *powershell) WriteCompletion(w io.Writer) error {
return generateCompletion(pshCompletionTemplate)(w, p.programName)
}

var _ Shell = &powershell{}

const pshCompletionTemplate = `
# === BEGIN {{.Name}} COMPLETION ===
# Escaping output sourced from:
# https://github.com/spf13/cobra/blob/e94f6d0dd9a5e5738dca6bce03c4b1207ffbc0ec/powershell_completions.go#L47
filter _{{.Name}}_escapeStringWithSpecialChars {
Expand Down Expand Up @@ -37,6 +81,6 @@ $_{{.Name}}_completions = {
}
rm env:COMPLETION_MODE
}
Register-ArgumentCompleter -CommandName {{.Name}} -ScriptBlock $_{{.Name}}_completions
# === END {{.Name}} COMPLETION ===
`
Loading

0 comments on commit 460dd41

Please # to comment.