1 package cobra |
1 package cobra |
2 |
2 |
3 import ( |
3 import ( |
4 "bytes" |
4 "encoding/json" |
5 "fmt" |
5 "fmt" |
6 "io" |
6 "io" |
7 "os" |
7 "os" |
|
8 "sort" |
8 "strings" |
9 "strings" |
|
10 "text/template" |
|
11 |
|
12 "github.com/spf13/pflag" |
9 ) |
13 ) |
|
14 |
|
15 const ( |
|
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 |
|
22 var ( |
|
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 .}} |
|
33 function {{$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" -}} |
|
63 function {{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. |
|
89 type zshCompArgsAnnotation map[int]zshCompArgHint |
|
90 |
|
91 type 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 } |
10 |
99 |
11 // GenZshCompletionFile generates zsh completion file. |
100 // GenZshCompletionFile generates zsh completion file. |
12 func (c *Command) GenZshCompletionFile(filename string) error { |
101 func (c *Command) GenZshCompletionFile(filename string) error { |
13 outFile, err := os.Create(filename) |
102 outFile, err := os.Create(filename) |
14 if err != nil { |
103 if err != nil { |
17 defer outFile.Close() |
106 defer outFile.Close() |
18 |
107 |
19 return c.GenZshCompletion(outFile) |
108 return c.GenZshCompletion(outFile) |
20 } |
109 } |
21 |
110 |
22 // GenZshCompletion generates a zsh completion file and writes to the passed writer. |
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. |
23 func (c *Command) GenZshCompletion(w io.Writer) error { |
114 func (c *Command) GenZshCompletion(w io.Writer) error { |
24 buf := new(bytes.Buffer) |
115 tmpl, err := template.New("Main").Funcs(zshCompFuncMap).Parse(zshCompletionText) |
25 |
116 if err != nil { |
26 writeHeader(buf, c) |
117 return fmt.Errorf("error creating zsh completion template: %v", err) |
27 maxDepth := maxDepth(c) |
118 } |
28 writeLevelMapping(buf, maxDepth) |
119 return tmpl.Execute(w, c.Root()) |
29 writeLevelCases(buf, maxDepth, c) |
120 } |
30 |
121 |
31 _, err := buf.WriteTo(w) |
122 // MarkZshCompPositionalArgumentFile marks the specified argument (first |
32 return err |
123 // argument is 1) as completed by file selection. patterns (e.g. "*.txt") are |
33 } |
124 // optional - if not provided the completion will search for all files. |
34 |
125 func (c *Command) MarkZshCompPositionalArgumentFile(argPosition int, patterns ...string) error { |
35 func writeHeader(w io.Writer, cmd *Command) { |
126 if argPosition < 1 { |
36 fmt.Fprintf(w, "#compdef %s\n\n", cmd.Name()) |
127 return fmt.Errorf("Invalid argument position (%d)", argPosition) |
37 } |
128 } |
38 |
129 annotation, err := c.zshCompGetArgsAnnotations() |
39 func maxDepth(c *Command) int { |
130 if err != nil { |
40 if len(c.Commands()) == 0 { |
131 return err |
41 return 0 |
132 } |
42 } |
133 if c.zshcompArgsAnnotationnIsDuplicatePosition(annotation, argPosition) { |
43 maxDepthSub := 0 |
134 return fmt.Errorf("Duplicate annotation for positional argument at index %d", argPosition) |
44 for _, s := range c.Commands() { |
135 } |
45 subDepth := maxDepth(s) |
136 annotation[argPosition] = zshCompArgHint{ |
46 if subDepth > maxDepthSub { |
137 Tipe: zshCompArgumentFilenameComp, |
47 maxDepthSub = subDepth |
138 Options: patterns, |
48 } |
139 } |
49 } |
140 return c.zshCompSetArgsAnnotations(annotation) |
50 return 1 + maxDepthSub |
141 } |
51 } |
142 |
52 |
143 // MarkZshCompPositionalArgumentWords marks the specified positional argument |
53 func writeLevelMapping(w io.Writer, numLevels int) { |
144 // (first argument is 1) as completed by the provided words. At east one word |
54 fmt.Fprintln(w, `_arguments \`) |
145 // must be provided, spaces within words will be offered completion with |
55 for i := 1; i <= numLevels; i++ { |
146 // "word\ word". |
56 fmt.Fprintf(w, ` '%d: :->level%d' \`, i, i) |
147 func (c *Command) MarkZshCompPositionalArgumentWords(argPosition int, words ...string) error { |
57 fmt.Fprintln(w) |
148 if argPosition < 1 { |
58 } |
149 return fmt.Errorf("Invalid argument position (%d)", argPosition) |
59 fmt.Fprintf(w, ` '%d: :%s'`, numLevels+1, "_files") |
150 } |
60 fmt.Fprintln(w) |
151 if len(words) == 0 { |
61 } |
152 return fmt.Errorf("Trying to set empty word list for positional argument %d", argPosition) |
62 |
153 } |
63 func writeLevelCases(w io.Writer, maxDepth int, root *Command) { |
154 annotation, err := c.zshCompGetArgsAnnotations() |
64 fmt.Fprintln(w, "case $state in") |
155 if err != nil { |
65 defer fmt.Fprintln(w, "esac") |
156 return err |
66 |
157 } |
67 for i := 1; i <= maxDepth; i++ { |
158 if c.zshcompArgsAnnotationnIsDuplicatePosition(annotation, argPosition) { |
68 fmt.Fprintf(w, " level%d)\n", i) |
159 return fmt.Errorf("Duplicate annotation for positional argument at index %d", argPosition) |
69 writeLevel(w, root, i) |
160 } |
70 fmt.Fprintln(w, " ;;") |
161 annotation[argPosition] = zshCompArgHint{ |
71 } |
162 Tipe: zshCompArgumentWordComp, |
72 fmt.Fprintln(w, " *)") |
163 Options: words, |
73 fmt.Fprintln(w, " _arguments '*: :_files'") |
164 } |
74 fmt.Fprintln(w, " ;;") |
165 return c.zshCompSetArgsAnnotations(annotation) |
75 } |
166 } |
76 |
167 |
77 func writeLevel(w io.Writer, root *Command, i int) { |
168 func zshCompExtractArgumentCompletionHintsForRendering(c *Command) ([]string, error) { |
78 fmt.Fprintf(w, " case $words[%d] in\n", i) |
169 var result []string |
79 defer fmt.Fprintln(w, " esac") |
170 annotation, err := c.zshCompGetArgsAnnotations() |
80 |
171 if err != nil { |
81 commands := filterByLevel(root, i) |
172 return nil, err |
82 byParent := groupByParent(commands) |
173 } |
83 |
174 for k, v := range annotation { |
84 for p, c := range byParent { |
175 s, err := zshCompRenderZshCompArgHint(k, v) |
85 names := names(c) |
176 if err != nil { |
86 fmt.Fprintf(w, " %s)\n", p) |
177 return nil, err |
87 fmt.Fprintf(w, " _arguments '%d: :(%s)'\n", i, strings.Join(names, " ")) |
178 } |
88 fmt.Fprintln(w, " ;;") |
179 result = append(result, s) |
89 } |
180 } |
90 fmt.Fprintln(w, " *)") |
181 if len(c.ValidArgs) > 0 { |
91 fmt.Fprintln(w, " _arguments '*: :_files'") |
182 if _, positionOneExists := annotation[1]; !positionOneExists { |
92 fmt.Fprintln(w, " ;;") |
183 s, err := zshCompRenderZshCompArgHint(1, zshCompArgHint{ |
93 |
184 Tipe: zshCompArgumentWordComp, |
94 } |
185 Options: c.ValidArgs, |
95 |
186 }) |
96 func filterByLevel(c *Command, l int) []*Command { |
187 if err != nil { |
97 cs := make([]*Command, 0) |
188 return nil, err |
98 if l == 0 { |
189 } |
99 cs = append(cs, c) |
190 result = append(result, s) |
100 return cs |
191 } |
101 } |
192 } |
102 for _, s := range c.Commands() { |
193 sort.Strings(result) |
103 cs = append(cs, filterByLevel(s, l-1)...) |
194 return result, nil |
104 } |
195 } |
105 return cs |
196 |
106 } |
197 func zshCompRenderZshCompArgHint(i int, z zshCompArgHint) (string, error) { |
107 |
198 switch t := z.Tipe; t { |
108 func groupByParent(commands []*Command) map[string][]*Command { |
199 case zshCompArgumentFilenameComp: |
109 m := make(map[string][]*Command) |
200 var globs []string |
110 for _, c := range commands { |
201 for _, g := range z.Options { |
111 parent := c.Parent() |
202 globs = append(globs, fmt.Sprintf(`-g "%s"`, g)) |
112 if parent == nil { |
203 } |
113 continue |
204 return fmt.Sprintf(`'%d: :_files %s'`, i, strings.Join(globs, " ")), nil |
114 } |
205 case zshCompArgumentWordComp: |
115 m[parent.Name()] = append(m[parent.Name()], c) |
206 var words []string |
116 } |
207 for _, w := range z.Options { |
117 return m |
208 words = append(words, fmt.Sprintf("%q", w)) |
118 } |
209 } |
119 |
210 return fmt.Sprintf(`'%d: :(%s)'`, i, strings.Join(words, " ")), nil |
120 func names(commands []*Command) []string { |
211 default: |
121 ns := make([]string, len(commands)) |
212 return "", fmt.Errorf("Invalid zsh argument completion annotation: %s", t) |
122 for i, c := range commands { |
213 } |
123 ns[i] = c.Name() |
214 } |
124 } |
215 |
125 return ns |
216 func (c *Command) zshcompArgsAnnotationnIsDuplicatePosition(annotation zshCompArgsAnnotation, position int) bool { |
126 } |
217 _, dup := annotation[position] |
|
218 return dup |
|
219 } |
|
220 |
|
221 func (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 |
|
234 func (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 |
|
246 func zshCompGenFuncName(c *Command) string { |
|
247 if c.HasParent() { |
|
248 return zshCompGenFuncName(c.Parent()) + "_" + c.Name() |
|
249 } |
|
250 return "_" + c.Name() |
|
251 } |
|
252 |
|
253 func 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. |
|
270 func zshCompGenFlagEntryForArguments(f *pflag.Flag) string { |
|
271 if f.Name == "" || f.Shorthand == "" { |
|
272 return zshCompGenFlagEntryForSingleOptionFlag(f) |
|
273 } |
|
274 return zshCompGenFlagEntryForMultiOptionFlag(f) |
|
275 } |
|
276 |
|
277 func 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 |
|
293 func 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 |
|
308 func 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 |
|
329 func zshCompFlagCouldBeSpecifiedMoreThenOnce(f *pflag.Flag) bool { |
|
330 return strings.Contains(f.Value.Type(), "Slice") || |
|
331 strings.Contains(f.Value.Type(), "Array") |
|
332 } |
|
333 |
|
334 func zshCompQuoteFlagDescription(s string) string { |
|
335 return strings.Replace(s, "'", `'\''`, -1) |
|
336 } |