1package cobra
  2
  3import (
  4	"encoding/json"
  5	"fmt"
  6	"io"
  7	"os"
  8	"sort"
  9	"strings"
 10	"text/template"
 11
 12	"github.com/spf13/pflag"
 13)
 14
 15const (
 16	zshCompArgumentAnnotation   = "cobra_annotations_zsh_completion_argument_annotation"
 17	zshCompArgumentFilenameComp = "cobra_annotations_zsh_completion_argument_file_completion"
 18	zshCompArgumentWordComp     = "cobra_annotations_zsh_completion_argument_word_completion"
 19	zshCompDirname              = "cobra_annotations_zsh_dirname"
 20)
 21
 22var (
 23	zshCompFuncMap = template.FuncMap{
 24		"genZshFuncName":              zshCompGenFuncName,
 25		"extractFlags":                zshCompExtractFlag,
 26		"genFlagEntryForZshArguments": zshCompGenFlagEntryForArguments,
 27		"extractArgsCompletions":      zshCompExtractArgumentCompletionHintsForRendering,
 28	}
 29	zshCompletionText = `
 30{{/* should accept Command (that contains subcommands) as parameter */}}
 31{{define "argumentsC" -}}
 32{{ $cmdPath := genZshFuncName .}}
 33function {{$cmdPath}} {
 34  local -a commands
 35
 36  _arguments -C \{{- range extractFlags .}}
 37    {{genFlagEntryForZshArguments .}} \{{- end}}
 38    "1: :->cmnds" \
 39    "*::arg:->args"
 40
 41  case $state in
 42  cmnds)
 43    commands=({{range .Commands}}{{if not .Hidden}}
 44      "{{.Name}}:{{.Short}}"{{end}}{{end}}
 45    )
 46    _describe "command" commands
 47    ;;
 48  esac
 49
 50  case "$words[1]" in {{- range .Commands}}{{if not .Hidden}}
 51  {{.Name}})
 52    {{$cmdPath}}_{{.Name}}
 53    ;;{{end}}{{end}}
 54  esac
 55}
 56{{range .Commands}}{{if not .Hidden}}
 57{{template "selectCmdTemplate" .}}
 58{{- end}}{{end}}
 59{{- end}}
 60
 61{{/* should accept Command without subcommands as parameter */}}
 62{{define "arguments" -}}
 63function {{genZshFuncName .}} {
 64{{"  _arguments"}}{{range extractFlags .}} \
 65    {{genFlagEntryForZshArguments . -}}
 66{{end}}{{range extractArgsCompletions .}} \
 67    {{.}}{{end}}
 68}
 69{{end}}
 70
 71{{/* dispatcher for commands with or without subcommands */}}
 72{{define "selectCmdTemplate" -}}
 73{{if .Hidden}}{{/* ignore hidden*/}}{{else -}}
 74{{if .Commands}}{{template "argumentsC" .}}{{else}}{{template "arguments" .}}{{end}}
 75{{- end}}
 76{{- end}}
 77
 78{{/* template entry point */}}
 79{{define "Main" -}}
 80#compdef _{{.Name}} {{.Name}}
 81
 82{{template "selectCmdTemplate" .}}
 83{{end}}
 84`
 85)
 86
 87// zshCompArgsAnnotation is used to encode/decode zsh completion for
 88// arguments to/from Command.Annotations.
 89type zshCompArgsAnnotation map[int]zshCompArgHint
 90
 91type zshCompArgHint struct {
 92	// Indicates the type of the completion to use. One of:
 93	// zshCompArgumentFilenameComp or zshCompArgumentWordComp
 94	Tipe string `json:"type"`
 95
 96	// A value for the type above (globs for file completion or words)
 97	Options []string `json:"options"`
 98}
 99
100// GenZshCompletionFile generates zsh completion file.
101func (c *Command) GenZshCompletionFile(filename string) error {
102	outFile, err := os.Create(filename)
103	if err != nil {
104		return err
105	}
106	defer outFile.Close()
107
108	return c.GenZshCompletion(outFile)
109}
110
111// GenZshCompletion generates a zsh completion file and writes to the passed
112// writer. The completion always run on the root command regardless of the
113// command it was called from.
114func (c *Command) GenZshCompletion(w io.Writer) error {
115	tmpl, err := template.New("Main").Funcs(zshCompFuncMap).Parse(zshCompletionText)
116	if err != nil {
117		return fmt.Errorf("error creating zsh completion template: %v", err)
118	}
119	return tmpl.Execute(w, c.Root())
120}
121
122// MarkZshCompPositionalArgumentFile marks the specified argument (first
123// argument is 1) as completed by file selection. patterns (e.g. "*.txt") are
124// optional - if not provided the completion will search for all files.
125func (c *Command) MarkZshCompPositionalArgumentFile(argPosition int, patterns ...string) error {
126	if argPosition < 1 {
127		return fmt.Errorf("Invalid argument position (%d)", argPosition)
128	}
129	annotation, err := c.zshCompGetArgsAnnotations()
130	if err != nil {
131		return err
132	}
133	if c.zshcompArgsAnnotationnIsDuplicatePosition(annotation, argPosition) {
134		return fmt.Errorf("Duplicate annotation for positional argument at index %d", argPosition)
135	}
136	annotation[argPosition] = zshCompArgHint{
137		Tipe:    zshCompArgumentFilenameComp,
138		Options: patterns,
139	}
140	return c.zshCompSetArgsAnnotations(annotation)
141}
142
143// MarkZshCompPositionalArgumentWords marks the specified positional argument
144// (first argument is 1) as completed by the provided words. At east one word
145// must be provided, spaces within words will be offered completion with
146// "word\ word".
147func (c *Command) MarkZshCompPositionalArgumentWords(argPosition int, words ...string) error {
148	if argPosition < 1 {
149		return fmt.Errorf("Invalid argument position (%d)", argPosition)
150	}
151	if len(words) == 0 {
152		return fmt.Errorf("Trying to set empty word list for positional argument %d", argPosition)
153	}
154	annotation, err := c.zshCompGetArgsAnnotations()
155	if err != nil {
156		return err
157	}
158	if c.zshcompArgsAnnotationnIsDuplicatePosition(annotation, argPosition) {
159		return fmt.Errorf("Duplicate annotation for positional argument at index %d", argPosition)
160	}
161	annotation[argPosition] = zshCompArgHint{
162		Tipe:    zshCompArgumentWordComp,
163		Options: words,
164	}
165	return c.zshCompSetArgsAnnotations(annotation)
166}
167
168func zshCompExtractArgumentCompletionHintsForRendering(c *Command) ([]string, error) {
169	var result []string
170	annotation, err := c.zshCompGetArgsAnnotations()
171	if err != nil {
172		return nil, err
173	}
174	for k, v := range annotation {
175		s, err := zshCompRenderZshCompArgHint(k, v)
176		if err != nil {
177			return nil, err
178		}
179		result = append(result, s)
180	}
181	if len(c.ValidArgs) > 0 {
182		if _, positionOneExists := annotation[1]; !positionOneExists {
183			s, err := zshCompRenderZshCompArgHint(1, zshCompArgHint{
184				Tipe:    zshCompArgumentWordComp,
185				Options: c.ValidArgs,
186			})
187			if err != nil {
188				return nil, err
189			}
190			result = append(result, s)
191		}
192	}
193	sort.Strings(result)
194	return result, nil
195}
196
197func zshCompRenderZshCompArgHint(i int, z zshCompArgHint) (string, error) {
198	switch t := z.Tipe; t {
199	case zshCompArgumentFilenameComp:
200		var globs []string
201		for _, g := range z.Options {
202			globs = append(globs, fmt.Sprintf(`-g "%s"`, g))
203		}
204		return fmt.Sprintf(`'%d: :_files %s'`, i, strings.Join(globs, " ")), nil
205	case zshCompArgumentWordComp:
206		var words []string
207		for _, w := range z.Options {
208			words = append(words, fmt.Sprintf("%q", w))
209		}
210		return fmt.Sprintf(`'%d: :(%s)'`, i, strings.Join(words, " ")), nil
211	default:
212		return "", fmt.Errorf("Invalid zsh argument completion annotation: %s", t)
213	}
214}
215
216func (c *Command) zshcompArgsAnnotationnIsDuplicatePosition(annotation zshCompArgsAnnotation, position int) bool {
217	_, dup := annotation[position]
218	return dup
219}
220
221func (c *Command) zshCompGetArgsAnnotations() (zshCompArgsAnnotation, error) {
222	annotation := make(zshCompArgsAnnotation)
223	annotationString, ok := c.Annotations[zshCompArgumentAnnotation]
224	if !ok {
225		return annotation, nil
226	}
227	err := json.Unmarshal([]byte(annotationString), &annotation)
228	if err != nil {
229		return annotation, fmt.Errorf("Error unmarshaling zsh argument annotation: %v", err)
230	}
231	return annotation, nil
232}
233
234func (c *Command) zshCompSetArgsAnnotations(annotation zshCompArgsAnnotation) error {
235	jsn, err := json.Marshal(annotation)
236	if err != nil {
237		return fmt.Errorf("Error marshaling zsh argument annotation: %v", err)
238	}
239	if c.Annotations == nil {
240		c.Annotations = make(map[string]string)
241	}
242	c.Annotations[zshCompArgumentAnnotation] = string(jsn)
243	return nil
244}
245
246func zshCompGenFuncName(c *Command) string {
247	if c.HasParent() {
248		return zshCompGenFuncName(c.Parent()) + "_" + c.Name()
249	}
250	return "_" + c.Name()
251}
252
253func zshCompExtractFlag(c *Command) []*pflag.Flag {
254	var flags []*pflag.Flag
255	c.LocalFlags().VisitAll(func(f *pflag.Flag) {
256		if !f.Hidden {
257			flags = append(flags, f)
258		}
259	})
260	c.InheritedFlags().VisitAll(func(f *pflag.Flag) {
261		if !f.Hidden {
262			flags = append(flags, f)
263		}
264	})
265	return flags
266}
267
268// zshCompGenFlagEntryForArguments returns an entry that matches _arguments
269// zsh-completion parameters. It's too complicated to generate in a template.
270func zshCompGenFlagEntryForArguments(f *pflag.Flag) string {
271	if f.Name == "" || f.Shorthand == "" {
272		return zshCompGenFlagEntryForSingleOptionFlag(f)
273	}
274	return zshCompGenFlagEntryForMultiOptionFlag(f)
275}
276
277func zshCompGenFlagEntryForSingleOptionFlag(f *pflag.Flag) string {
278	var option, multiMark, extras string
279
280	if zshCompFlagCouldBeSpecifiedMoreThenOnce(f) {
281		multiMark = "*"
282	}
283
284	option = "--" + f.Name
285	if option == "--" {
286		option = "-" + f.Shorthand
287	}
288	extras = zshCompGenFlagEntryExtras(f)
289
290	return fmt.Sprintf(`'%s%s[%s]%s'`, multiMark, option, zshCompQuoteFlagDescription(f.Usage), extras)
291}
292
293func zshCompGenFlagEntryForMultiOptionFlag(f *pflag.Flag) string {
294	var options, parenMultiMark, curlyMultiMark, extras string
295
296	if zshCompFlagCouldBeSpecifiedMoreThenOnce(f) {
297		parenMultiMark = "*"
298		curlyMultiMark = "\\*"
299	}
300
301	options = fmt.Sprintf(`'(%s-%s %s--%s)'{%s-%s,%s--%s}`,
302		parenMultiMark, f.Shorthand, parenMultiMark, f.Name, curlyMultiMark, f.Shorthand, curlyMultiMark, f.Name)
303	extras = zshCompGenFlagEntryExtras(f)
304
305	return fmt.Sprintf(`%s'[%s]%s'`, options, zshCompQuoteFlagDescription(f.Usage), extras)
306}
307
308func zshCompGenFlagEntryExtras(f *pflag.Flag) string {
309	if f.NoOptDefVal != "" {
310		return ""
311	}
312
313	extras := ":" // allow options for flag (even without assistance)
314	for key, values := range f.Annotations {
315		switch key {
316		case zshCompDirname:
317			extras = fmt.Sprintf(":filename:_files -g %q", values[0])
318		case BashCompFilenameExt:
319			extras = ":filename:_files"
320			for _, pattern := range values {
321				extras = extras + fmt.Sprintf(` -g "%s"`, pattern)
322			}
323		}
324	}
325
326	return extras
327}
328
329func zshCompFlagCouldBeSpecifiedMoreThenOnce(f *pflag.Flag) bool {
330	return strings.Contains(f.Value.Type(), "Slice") ||
331		strings.Contains(f.Value.Type(), "Array")
332}
333
334func zshCompQuoteFlagDescription(s string) string {
335	return strings.Replace(s, "'", `'\''`, -1)
336}