diff --git a/cmd/netssh/main.go b/cmd/netssh/main.go index 7d52e17..d2e005e 100644 --- a/cmd/netssh/main.go +++ b/cmd/netssh/main.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "io" "os" "os/exec" "path/filepath" @@ -311,8 +312,13 @@ func rootCmd() *cobra.Command { c := cache.New(cfg.CachePath(), cfg.Cache.TTL) _ = c.Load() - var completions []cobra.Completion - for _, e := range c.Search(toComplete) { + hosts := c.Search(toComplete) + sort.Slice(hosts, func(i, j int) bool { return hosts[i].Name < hosts[j].Name }) + if limit := cfg.Completion.MaxResults; limit > 0 && len(hosts) > limit { + hosts = hosts[:limit] + } + completions := make([]cobra.Completion, 0, len(hosts)) + for _, e := range hosts { completions = append(completions, cobra.Completion(e.Name)) } // Subcommands at the end, after all hostnames. @@ -332,6 +338,17 @@ func rootCmd() *cobra.Command { for _, cmd := range root.Commands() { if cmd.Name() == "completion" { cmd.AddCommand(completionInstallCmd(root)) + // Cobra's zsh script spawns a subshell per line and sends every host + // through _describe, which hangs with thousands of cached hosts. + for _, sub := range cmd.Commands() { + if sub.Name() == "zsh" { + sub.RunE = func(cmd *cobra.Command, args []string) error { + _, err := io.WriteString(cmd.OutOrStdout(), hook.ZshCompletionScript) + return err + } + sub.Run = nil + } + } break } } @@ -369,12 +386,8 @@ func completionInstallCmd(root *cobra.Command) *cobra.Command { case "zsh": dir = filepath.Join(os.Getenv("HOME"), ".zfunc") file = filepath.Join(dir, "_netssh") - gen = func() ([]byte, error) { - var buf strings.Builder - err := root.GenZshCompletion(&buf) - return []byte(buf.String()), err - } - note = "Make sure ~/.zfunc is in your fpath:\n fpath=(~/.zfunc $fpath)\n autoload -Uz compinit && compinit" + gen = func() ([]byte, error) { return []byte(hook.ZshCompletionScript), nil } + note = hook.ZshCompletionNote(dir) case "fish": dir = filepath.Join(hook.FishConfigDir(), "completions") file = filepath.Join(dir, "netssh.fish") diff --git a/internal/config/config.go b/internal/config/config.go index 6c686c0..a4377ad 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -11,13 +11,14 @@ import ( ) type Config struct { - Shell ShellConfig `mapstructure:"shell"` - NetBox NetBoxConfig `mapstructure:"netbox"` - NetBoxes []NetBoxConfig `mapstructure:"netboxes"` - Resolver ResolverConfig `mapstructure:"resolver"` - Cache CacheConfig `mapstructure:"cache"` - SSH SSHConfig `mapstructure:"ssh"` - Shortcuts ShortcutsConfig `mapstructure:"shortcuts"` + Shell ShellConfig `mapstructure:"shell"` + NetBox NetBoxConfig `mapstructure:"netbox"` + NetBoxes []NetBoxConfig `mapstructure:"netboxes"` + Resolver ResolverConfig `mapstructure:"resolver"` + Cache CacheConfig `mapstructure:"cache"` + SSH SSHConfig `mapstructure:"ssh"` + Shortcuts ShortcutsConfig `mapstructure:"shortcuts"` + Completion CompletionConfig `mapstructure:"completion"` } type ShellConfig struct { @@ -47,6 +48,11 @@ type SSHConfig struct { DefaultPort int `mapstructure:"default_port"` } +type CompletionConfig struct { + // MaxResults limits the number of hostnames offered per completion; 0 = unlimited. + MaxResults int `mapstructure:"max_results"` +} + type ShortcutsConfig struct { Domains []string `mapstructure:"domains"` StripHyphens bool `mapstructure:"strip_hyphens"` diff --git a/internal/hook/completion.go b/internal/hook/completion.go index 1f42980..11d1495 100644 --- a/internal/hook/completion.go +++ b/internal/hook/completion.go @@ -83,7 +83,7 @@ func commandCompletionPath(shell, alias string) string { case "bash": return filepath.Join(os.Getenv("HOME"), ".local", "share", "bash-completion", "completions", alias) case "zsh": - return filepath.Join(os.Getenv("HOME"), ".zfunc", "_"+alias) + return filepath.Join(ZshFuncDir(), "_"+alias) default: return filepath.Join(FishConfigDir(), "completions", alias+".fish") } diff --git a/internal/hook/completion_test.go b/internal/hook/completion_test.go index 215175b..c5f3f14 100644 --- a/internal/hook/completion_test.go +++ b/internal/hook/completion_test.go @@ -88,7 +88,7 @@ func TestCommandAliasCompletesServerNamesAcrossShells(t *testing.T) { case "bash": err = root.GenBashCompletionV2(&canonical, true) case "zsh": - err = root.GenZshCompletion(&canonical) + _, err = canonical.WriteString(hook.ZshCompletionScript) case "fish": err = root.GenFishCompletion(&canonical, true) } @@ -149,7 +149,8 @@ compinit -u -D alias j=netssh autoload -Uz _j # Capture matches without requiring an interactive ZLE completion widget. -_describe() { print -rl -- "${completions[@]}"; } +_describe() { print -rl -- "${(P)4}"; } +compadd() { print -rl -- "$@"; } words=(j ser) CURRENT=2 _j @@ -169,3 +170,38 @@ complete --do-complete 'j ser' }) } } + +func TestZshCompletionSplitsHostsAndCommands(t *testing.T) { + zsh, err := exec.LookPath("zsh") + if err != nil { + t.Skip("zsh unavailable") + } + dir := t.TempDir() + backend := "#!/bin/sh\nprintf 'web01\\nweb:02\\ncache\\tManage the cache\\n:36\\n'\n" + if err := os.WriteFile(filepath.Join(dir, "netssh"), []byte(backend), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "_netssh"), []byte(hook.ZshCompletionScript), 0o644); err != nil { + t.Fatal(err) + } + script := `compdef() { :; } +source "$DIR/_netssh" +_describe() { print -r -- "describe ${(P)4}"; } +compadd() { print -r -- "compadd $*"; } +_files() { print -r -- files; } +words=(netssh "") +CURRENT=2 +_netssh +` + cmd := exec.Command(zsh, "-f", "-c", script) + cmd.Env = append(os.Environ(), "DIR="+dir, "PATH="+dir+string(os.PathListSeparator)+os.Getenv("PATH")) + out, err := cmd.CombinedOutput() + if err != nil { + t.Fatalf("%v\n%s", err, out) + } + // Directive 36 = NoFileComp|KeepOrder: unsorted host group, no file fallback. + want := "describe cache:Manage the cache\ncompadd -V hosts -X %Bhosts%b -- web01 web:02\n" + if string(out) != want { + t.Fatalf("got:\n%s\nwant:\n%s", out, want) + } +} diff --git a/internal/hook/zsh.go b/internal/hook/zsh.go new file mode 100644 index 0000000..a9db706 --- /dev/null +++ b/internal/hook/zsh.go @@ -0,0 +1,99 @@ +package hook + +import ( + "fmt" + "os" + "path/filepath" +) + +// ZshCompletionScript replaces Cobra's generated zsh completion. Cobra's script +// spawns a subshell per candidate and passes every host through _describe, +// which takes tens of seconds for thousands of hosts. This script splits the +// candidates with parameter expansion only: entries with a description +// (subcommands, flags) go through _describe, plain hostnames through compadd. +// +// The file works both from fpath (autoloaded as _netssh) and when sourced. +const ZshCompletionScript = `#compdef netssh +compdef _netssh netssh + +# netssh zsh completion — generated by netssh, do not edit. +_netssh() { + local -a lines described plain + local out directive=0 entry ret=1 + local -a nospace + + out=$(netssh __complete "${(@)words[2,CURRENT]}" 2>/dev/null) || return 1 + lines=("${(@f)out}") + if [[ ${lines[-1]} == :<-> ]]; then + directive=${lines[-1]#:} + lines[-1]=() + fi + # Directive 1: error. + (( directive & 1 )) && return 1 + lines=(${lines:#_activeHelp_*}) + lines=(${lines:#}) + # Directive 2: no trailing space. + (( directive & 2 )) && nospace=(-S '') + + # Directive 8: the lines are file extensions to filter by. + if (( directive & 8 )); then + local -a globs + for entry in $lines; do globs+=("*.${entry}"); done + _files -g "(${(j:|:)globs})" + return + fi + # Directive 16: complete directories only. + if (( directive & 16 )); then + _files -/ + return + fi + + described=(${(M)lines:#*$'\t'*}) + plain=(${lines:#*$'\t'*}) + if (( $#described )); then + local -a items + for entry in $described; do + items+=("${${entry%%$'\t'*}//:/\\:}:${entry#*$'\t'}") + done + _describe -t commands 'netssh command' items $nospace && ret=0 + fi + if (( $#plain )); then + # Directive 32: keep the order supplied by netssh. + if (( directive & 32 )); then + compadd -V hosts -X '%Bhosts%b' $nospace -- $plain && ret=0 + else + compadd -J hosts -X '%Bhosts%b' $nospace -- $plain && ret=0 + fi + fi + # Directive 4: no file completion as a fallback. + if (( ! $#lines && ! (directive & 4) )); then + _files && ret=0 + fi + return ret +} + +if [ "$funcstack[1]" = "_netssh" ]; then + _netssh "$@" +fi +` + +// ZshCompletionNote tells the user how zsh picks up a completion file in dir. +func ZshCompletionNote(dir string) string { + dump := filepath.Join(zshDir(), ".zcompdump") + return fmt.Sprintf(`zsh loads it from fpath. The netssh block in your .zshrc adds %s +(netssh hook install). Otherwise add this before compinit: + fpath=(%s $fpath) +Then refresh the completion cache: rm -f %s* && exec zsh`, dir, dir, dump) +} + +func zshDir() string { + if dir := os.Getenv("ZDOTDIR"); dir != "" { + return dir + } + return os.Getenv("HOME") +} + +// ZshFuncDir is the user directory netssh installs zsh completion files into. +func ZshFuncDir() string { + return filepath.Join(os.Getenv("HOME"), ".zfunc") +}