337 lines
		
	
	
		
			9.6 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			337 lines
		
	
	
		
			9.6 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
package cobra
 | 
						|
 | 
						|
import (
 | 
						|
	"encoding/json"
 | 
						|
	"fmt"
 | 
						|
	"io"
 | 
						|
	"os"
 | 
						|
	"sort"
 | 
						|
	"strings"
 | 
						|
	"text/template"
 | 
						|
 | 
						|
	"github.com/spf13/pflag"
 | 
						|
)
 | 
						|
 | 
						|
const (
 | 
						|
	zshCompArgumentAnnotation   = "cobra_annotations_zsh_completion_argument_annotation"
 | 
						|
	zshCompArgumentFilenameComp = "cobra_annotations_zsh_completion_argument_file_completion"
 | 
						|
	zshCompArgumentWordComp     = "cobra_annotations_zsh_completion_argument_word_completion"
 | 
						|
	zshCompDirname              = "cobra_annotations_zsh_dirname"
 | 
						|
)
 | 
						|
 | 
						|
var (
 | 
						|
	zshCompFuncMap = template.FuncMap{
 | 
						|
		"genZshFuncName":              zshCompGenFuncName,
 | 
						|
		"extractFlags":                zshCompExtractFlag,
 | 
						|
		"genFlagEntryForZshArguments": zshCompGenFlagEntryForArguments,
 | 
						|
		"extractArgsCompletions":      zshCompExtractArgumentCompletionHintsForRendering,
 | 
						|
	}
 | 
						|
	zshCompletionText = `
 | 
						|
{{/* should accept Command (that contains subcommands) as parameter */}}
 | 
						|
{{define "argumentsC" -}}
 | 
						|
{{ $cmdPath := genZshFuncName .}}
 | 
						|
function {{$cmdPath}} {
 | 
						|
  local -a commands
 | 
						|
 | 
						|
  _arguments -C \{{- range extractFlags .}}
 | 
						|
    {{genFlagEntryForZshArguments .}} \{{- end}}
 | 
						|
    "1: :->cmnds" \
 | 
						|
    "*::arg:->args"
 | 
						|
 | 
						|
  case $state in
 | 
						|
  cmnds)
 | 
						|
    commands=({{range .Commands}}{{if not .Hidden}}
 | 
						|
      "{{.Name}}:{{.Short}}"{{end}}{{end}}
 | 
						|
    )
 | 
						|
    _describe "command" commands
 | 
						|
    ;;
 | 
						|
  esac
 | 
						|
 | 
						|
  case "$words[1]" in {{- range .Commands}}{{if not .Hidden}}
 | 
						|
  {{.Name}})
 | 
						|
    {{$cmdPath}}_{{.Name}}
 | 
						|
    ;;{{end}}{{end}}
 | 
						|
  esac
 | 
						|
}
 | 
						|
{{range .Commands}}{{if not .Hidden}}
 | 
						|
{{template "selectCmdTemplate" .}}
 | 
						|
{{- end}}{{end}}
 | 
						|
{{- end}}
 | 
						|
 | 
						|
{{/* should accept Command without subcommands as parameter */}}
 | 
						|
{{define "arguments" -}}
 | 
						|
function {{genZshFuncName .}} {
 | 
						|
{{"  _arguments"}}{{range extractFlags .}} \
 | 
						|
    {{genFlagEntryForZshArguments . -}}
 | 
						|
{{end}}{{range extractArgsCompletions .}} \
 | 
						|
    {{.}}{{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 _{{.Name}} {{.Name}}
 | 
						|
 | 
						|
{{template "selectCmdTemplate" .}}
 | 
						|
{{end}}
 | 
						|
`
 | 
						|
)
 | 
						|
 | 
						|
// zshCompArgsAnnotation is used to encode/decode zsh completion for
 | 
						|
// arguments to/from Command.Annotations.
 | 
						|
type zshCompArgsAnnotation map[int]zshCompArgHint
 | 
						|
 | 
						|
type zshCompArgHint struct {
 | 
						|
	// Indicates the type of the completion to use. One of:
 | 
						|
	// zshCompArgumentFilenameComp or zshCompArgumentWordComp
 | 
						|
	Tipe string `json:"type"`
 | 
						|
 | 
						|
	// A value for the type above (globs for file completion or words)
 | 
						|
	Options []string `json:"options"`
 | 
						|
}
 | 
						|
 | 
						|
// 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. The completion always run on the root command regardless of the
 | 
						|
// command it was called from.
 | 
						|
func (c *Command) GenZshCompletion(w io.Writer) error {
 | 
						|
	tmpl, err := template.New("Main").Funcs(zshCompFuncMap).Parse(zshCompletionText)
 | 
						|
	if err != nil {
 | 
						|
		return fmt.Errorf("error creating zsh completion template: %v", err)
 | 
						|
	}
 | 
						|
	return tmpl.Execute(w, c.Root())
 | 
						|
}
 | 
						|
 | 
						|
// MarkZshCompPositionalArgumentFile marks the specified argument (first
 | 
						|
// argument is 1) as completed by file selection. patterns (e.g. "*.txt") are
 | 
						|
// optional - if not provided the completion will search for all files.
 | 
						|
func (c *Command) MarkZshCompPositionalArgumentFile(argPosition int, patterns ...string) error {
 | 
						|
	if argPosition < 1 {
 | 
						|
		return fmt.Errorf("Invalid argument position (%d)", argPosition)
 | 
						|
	}
 | 
						|
	annotation, err := c.zshCompGetArgsAnnotations()
 | 
						|
	if err != nil {
 | 
						|
		return err
 | 
						|
	}
 | 
						|
	if c.zshcompArgsAnnotationnIsDuplicatePosition(annotation, argPosition) {
 | 
						|
		return fmt.Errorf("Duplicate annotation for positional argument at index %d", argPosition)
 | 
						|
	}
 | 
						|
	annotation[argPosition] = zshCompArgHint{
 | 
						|
		Tipe:    zshCompArgumentFilenameComp,
 | 
						|
		Options: patterns,
 | 
						|
	}
 | 
						|
	return c.zshCompSetArgsAnnotations(annotation)
 | 
						|
}
 | 
						|
 | 
						|
// MarkZshCompPositionalArgumentWords marks the specified positional argument
 | 
						|
// (first argument is 1) as completed by the provided words. At east one word
 | 
						|
// must be provided, spaces within words will be offered completion with
 | 
						|
// "word\ word".
 | 
						|
func (c *Command) MarkZshCompPositionalArgumentWords(argPosition int, words ...string) error {
 | 
						|
	if argPosition < 1 {
 | 
						|
		return fmt.Errorf("Invalid argument position (%d)", argPosition)
 | 
						|
	}
 | 
						|
	if len(words) == 0 {
 | 
						|
		return fmt.Errorf("Trying to set empty word list for positional argument %d", argPosition)
 | 
						|
	}
 | 
						|
	annotation, err := c.zshCompGetArgsAnnotations()
 | 
						|
	if err != nil {
 | 
						|
		return err
 | 
						|
	}
 | 
						|
	if c.zshcompArgsAnnotationnIsDuplicatePosition(annotation, argPosition) {
 | 
						|
		return fmt.Errorf("Duplicate annotation for positional argument at index %d", argPosition)
 | 
						|
	}
 | 
						|
	annotation[argPosition] = zshCompArgHint{
 | 
						|
		Tipe:    zshCompArgumentWordComp,
 | 
						|
		Options: words,
 | 
						|
	}
 | 
						|
	return c.zshCompSetArgsAnnotations(annotation)
 | 
						|
}
 | 
						|
 | 
						|
func zshCompExtractArgumentCompletionHintsForRendering(c *Command) ([]string, error) {
 | 
						|
	var result []string
 | 
						|
	annotation, err := c.zshCompGetArgsAnnotations()
 | 
						|
	if err != nil {
 | 
						|
		return nil, err
 | 
						|
	}
 | 
						|
	for k, v := range annotation {
 | 
						|
		s, err := zshCompRenderZshCompArgHint(k, v)
 | 
						|
		if err != nil {
 | 
						|
			return nil, err
 | 
						|
		}
 | 
						|
		result = append(result, s)
 | 
						|
	}
 | 
						|
	if len(c.ValidArgs) > 0 {
 | 
						|
		if _, positionOneExists := annotation[1]; !positionOneExists {
 | 
						|
			s, err := zshCompRenderZshCompArgHint(1, zshCompArgHint{
 | 
						|
				Tipe:    zshCompArgumentWordComp,
 | 
						|
				Options: c.ValidArgs,
 | 
						|
			})
 | 
						|
			if err != nil {
 | 
						|
				return nil, err
 | 
						|
			}
 | 
						|
			result = append(result, s)
 | 
						|
		}
 | 
						|
	}
 | 
						|
	sort.Strings(result)
 | 
						|
	return result, nil
 | 
						|
}
 | 
						|
 | 
						|
func zshCompRenderZshCompArgHint(i int, z zshCompArgHint) (string, error) {
 | 
						|
	switch t := z.Tipe; t {
 | 
						|
	case zshCompArgumentFilenameComp:
 | 
						|
		var globs []string
 | 
						|
		for _, g := range z.Options {
 | 
						|
			globs = append(globs, fmt.Sprintf(`-g "%s"`, g))
 | 
						|
		}
 | 
						|
		return fmt.Sprintf(`'%d: :_files %s'`, i, strings.Join(globs, " ")), nil
 | 
						|
	case zshCompArgumentWordComp:
 | 
						|
		var words []string
 | 
						|
		for _, w := range z.Options {
 | 
						|
			words = append(words, fmt.Sprintf("%q", w))
 | 
						|
		}
 | 
						|
		return fmt.Sprintf(`'%d: :(%s)'`, i, strings.Join(words, " ")), nil
 | 
						|
	default:
 | 
						|
		return "", fmt.Errorf("Invalid zsh argument completion annotation: %s", t)
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
func (c *Command) zshcompArgsAnnotationnIsDuplicatePosition(annotation zshCompArgsAnnotation, position int) bool {
 | 
						|
	_, dup := annotation[position]
 | 
						|
	return dup
 | 
						|
}
 | 
						|
 | 
						|
func (c *Command) zshCompGetArgsAnnotations() (zshCompArgsAnnotation, error) {
 | 
						|
	annotation := make(zshCompArgsAnnotation)
 | 
						|
	annotationString, ok := c.Annotations[zshCompArgumentAnnotation]
 | 
						|
	if !ok {
 | 
						|
		return annotation, nil
 | 
						|
	}
 | 
						|
	err := json.Unmarshal([]byte(annotationString), &annotation)
 | 
						|
	if err != nil {
 | 
						|
		return annotation, fmt.Errorf("Error unmarshaling zsh argument annotation: %v", err)
 | 
						|
	}
 | 
						|
	return annotation, nil
 | 
						|
}
 | 
						|
 | 
						|
func (c *Command) zshCompSetArgsAnnotations(annotation zshCompArgsAnnotation) error {
 | 
						|
	jsn, err := json.Marshal(annotation)
 | 
						|
	if err != nil {
 | 
						|
		return fmt.Errorf("Error marshaling zsh argument annotation: %v", err)
 | 
						|
	}
 | 
						|
	if c.Annotations == nil {
 | 
						|
		c.Annotations = make(map[string]string)
 | 
						|
	}
 | 
						|
	c.Annotations[zshCompArgumentAnnotation] = string(jsn)
 | 
						|
	return nil
 | 
						|
}
 | 
						|
 | 
						|
func zshCompGenFuncName(c *Command) string {
 | 
						|
	if c.HasParent() {
 | 
						|
		return zshCompGenFuncName(c.Parent()) + "_" + c.Name()
 | 
						|
	}
 | 
						|
	return "_" + c.Name()
 | 
						|
}
 | 
						|
 | 
						|
func zshCompExtractFlag(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
 | 
						|
}
 | 
						|
 | 
						|
// zshCompGenFlagEntryForArguments returns an entry that matches _arguments
 | 
						|
// zsh-completion parameters. It's too complicated to generate in a template.
 | 
						|
func zshCompGenFlagEntryForArguments(f *pflag.Flag) string {
 | 
						|
	if f.Name == "" || f.Shorthand == "" {
 | 
						|
		return zshCompGenFlagEntryForSingleOptionFlag(f)
 | 
						|
	}
 | 
						|
	return zshCompGenFlagEntryForMultiOptionFlag(f)
 | 
						|
}
 | 
						|
 | 
						|
func zshCompGenFlagEntryForSingleOptionFlag(f *pflag.Flag) string {
 | 
						|
	var option, multiMark, extras string
 | 
						|
 | 
						|
	if zshCompFlagCouldBeSpecifiedMoreThenOnce(f) {
 | 
						|
		multiMark = "*"
 | 
						|
	}
 | 
						|
 | 
						|
	option = "--" + f.Name
 | 
						|
	if option == "--" {
 | 
						|
		option = "-" + f.Shorthand
 | 
						|
	}
 | 
						|
	extras = zshCompGenFlagEntryExtras(f)
 | 
						|
 | 
						|
	return fmt.Sprintf(`'%s%s[%s]%s'`, multiMark, option, zshCompQuoteFlagDescription(f.Usage), extras)
 | 
						|
}
 | 
						|
 | 
						|
func zshCompGenFlagEntryForMultiOptionFlag(f *pflag.Flag) string {
 | 
						|
	var options, parenMultiMark, curlyMultiMark, extras string
 | 
						|
 | 
						|
	if zshCompFlagCouldBeSpecifiedMoreThenOnce(f) {
 | 
						|
		parenMultiMark = "*"
 | 
						|
		curlyMultiMark = "\\*"
 | 
						|
	}
 | 
						|
 | 
						|
	options = fmt.Sprintf(`'(%s-%s %s--%s)'{%s-%s,%s--%s}`,
 | 
						|
		parenMultiMark, f.Shorthand, parenMultiMark, f.Name, curlyMultiMark, f.Shorthand, curlyMultiMark, f.Name)
 | 
						|
	extras = zshCompGenFlagEntryExtras(f)
 | 
						|
 | 
						|
	return fmt.Sprintf(`%s'[%s]%s'`, options, zshCompQuoteFlagDescription(f.Usage), extras)
 | 
						|
}
 | 
						|
 | 
						|
func zshCompGenFlagEntryExtras(f *pflag.Flag) string {
 | 
						|
	if f.NoOptDefVal != "" {
 | 
						|
		return ""
 | 
						|
	}
 | 
						|
 | 
						|
	extras := ":" // allow options for flag (even without assistance)
 | 
						|
	for key, values := range f.Annotations {
 | 
						|
		switch key {
 | 
						|
		case zshCompDirname:
 | 
						|
			extras = fmt.Sprintf(":filename:_files -g %q", values[0])
 | 
						|
		case BashCompFilenameExt:
 | 
						|
			extras = ":filename:_files"
 | 
						|
			for _, pattern := range values {
 | 
						|
				extras = extras + fmt.Sprintf(` -g "%s"`, pattern)
 | 
						|
			}
 | 
						|
		}
 | 
						|
	}
 | 
						|
 | 
						|
	return extras
 | 
						|
}
 | 
						|
 | 
						|
func zshCompFlagCouldBeSpecifiedMoreThenOnce(f *pflag.Flag) bool {
 | 
						|
	return strings.Contains(f.Value.Type(), "Slice") ||
 | 
						|
		strings.Contains(f.Value.Type(), "Array")
 | 
						|
}
 | 
						|
 | 
						|
func zshCompQuoteFlagDescription(s string) string {
 | 
						|
	return strings.Replace(s, "'", `'\''`, -1)
 | 
						|
}
 |