zsh_completions.go - external/github.com/spf13/cobra - Git at Google

 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)
 }
	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)
	}