cobra/zsh_completions.go

package cobra

import (
	"fmt"
	"io"
	"os"
	"strings"
	"text/template"

	"github.com/spf13/pflag"
)

var (
	funcMap = template.FuncMap{
		"constructPath": constructPath,
		"subCmdList":    subCmdList,
		"extractFlags":  extractFlags,
		"cmdName":       cmdName,
		"simpleFlag":    simpleFlag,
	}
	zshCompletionText = `
{{/* for pflag.Flag (specifically annotations) */}}
{{define "flagAnnotations" -}}
{{with index .Annotations "cobra_annotation_bash_completion_filename_extensions"}}:filename:_files{{end}}
{{- end}}

{{/* for pflag.Flag with short and long options */}}
{{define "complexFlag" -}}
"(-{{.Shorthand}} --{{.Name}})"{-{{.Shorthand}},--{{.Name}}}"[{{.Usage}}]{{template "flagAnnotations" .}}"
{{- end}}

{{/* for pflag.Flag with either short or long options */}}
{{define "simpleFlag" -}}
"{{with .Name}}--{{.}}{{else}}-{{.Shorthand}}{{end}}[{{.Usage}}]{{template "flagAnnotations" .}}"
{{- end}}

{{/* should accept Command (that contains subcommands) as parameter */}}
{{define "argumentsC" -}}
function {{constructPath .}} {
  local -a commands

  _arguments -C \{{- range extractFlags .}}
    {{if simpleFlag .}}{{template "simpleFlag" .}}{{else}}{{template "complexFlag" .}}{{end}} \{{- end}}
    "1: :->cmnds" \
    "*::arg:->args"

  case $state in
  cmnds)
    commands=({{range .Commands}}{{if not .Hidden}}
      "{{cmdName .}}:{{.Short}}"{{end}}{{end}}
    )
    _describe "command" commands
    ;;
  esac

  case "$words[1]" in {{- range .Commands}}{{if not .Hidden}}
  {{cmdName .}})
    {{constructPath .}}
    ;;{{end}}{{end}}
  esac
}
{{range .Commands}}{{if not .Hidden}}
{{template "selectCmdTemplate" .}}
{{- end}}{{end}}
{{- end}}

{{/* should accept Command without subcommands as parameter */}}
{{define "arguments" -}}
function {{constructPath .}} {
{{"  _arguments"}}{{range extractFlags .}} \
    {{if simpleFlag .}}{{template "simpleFlag" .}}{{else}}{{template "complexFlag" .}}{{end -}}
{{end}}
}
{{end}}

{{/* dispatcher for commands with or without subcommands */}}
{{define "selectCmdTemplate" -}}
{{if .Hidden}}{{/* ignore hidden*/}}{{else -}}
{{if .Commands}}{{template "argumentsC" .}}{{else}}{{template "arguments" .}}{{end}}
{{- end}}
{{- end}}

{{/* template entry point */}}
{{define "Main" -}}
#compdef _{{cmdName .}} {{cmdName .}}

{{template "selectCmdTemplate" .}}
{{end}}
`
)

// GenZshCompletionFile generates zsh completion file.
func (c *Command) GenZshCompletionFile(filename string) error {
	outFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer outFile.Close()

	return c.GenZshCompletion(outFile)
}

// GenZshCompletion generates a zsh completion file and writes to the passed writer.
func (c *Command) GenZshCompletion(w io.Writer) error {
	tmpl, err := template.New("Main").Funcs(funcMap).Parse(zshCompletionText)
	if err != nil {
		return fmt.Errorf("error creating zsh completion template: %v", err)
	}
	return tmpl.Execute(w, c)
}

func constructPath(c *Command) string {
	var path []string
	tmpCmd := c
	path = append(path, tmpCmd.Name())

	for {
		if !tmpCmd.HasParent() {
			break
		}
		tmpCmd = tmpCmd.Parent()
		path = append(path, tmpCmd.Name())
	}

	// reverse path
	for left, right := 0, len(path)-1; left < right; left, right = left+1, right-1 {
		path[left], path[right] = path[right], path[left]
	}

	return "_" + strings.Join(path, "_")
}

// subCmdList returns a space separated list of subcommands names
func subCmdList(c *Command) string {
	var subCmds []string

	for _, cmd := range c.Commands() {
		if cmd.Hidden {
			continue
		}
		subCmds = append(subCmds, cmd.Name())
	}

	return strings.Join(subCmds, " ")
}

func extractFlags(c *Command) []*pflag.Flag {
	var flags []*pflag.Flag
	c.LocalFlags().VisitAll(func(f *pflag.Flag) {
		if !f.Hidden {
			flags = append(flags, f)
		}
	})
	c.InheritedFlags().VisitAll(func(f *pflag.Flag) {
		if !f.Hidden {
			flags = append(flags, f)
		}
	})
	return flags
}

// cmdName returns the command's innvocation
func cmdName(c *Command) string {
	return c.Name()
}

func simpleFlag(p *pflag.Flag) bool {
	return p.Name == "" || p.Shorthand == ""
}