zsh_completions.go 9.6 KB

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