bash_completions.go 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584
  1. package cobra
  2. import (
  3. "bytes"
  4. "fmt"
  5. "io"
  6. "os"
  7. "sort"
  8. "strings"
  9. "github.com/spf13/pflag"
  10. )
  11. // Annotations for Bash completion.
  12. const (
  13. BashCompFilenameExt = "cobra_annotation_bash_completion_filename_extensions"
  14. BashCompCustom = "cobra_annotation_bash_completion_custom"
  15. BashCompOneRequiredFlag = "cobra_annotation_bash_completion_one_required_flag"
  16. BashCompSubdirsInDir = "cobra_annotation_bash_completion_subdirs_in_dir"
  17. )
  18. func writePreamble(buf *bytes.Buffer, name string) {
  19. buf.WriteString(fmt.Sprintf("# bash completion for %-36s -*- shell-script -*-\n", name))
  20. buf.WriteString(fmt.Sprintf(`
  21. __%[1]s_debug()
  22. {
  23. if [[ -n ${BASH_COMP_DEBUG_FILE} ]]; then
  24. echo "$*" >> "${BASH_COMP_DEBUG_FILE}"
  25. fi
  26. }
  27. # Homebrew on Macs have version 1.3 of bash-completion which doesn't include
  28. # _init_completion. This is a very minimal version of that function.
  29. __%[1]s_init_completion()
  30. {
  31. COMPREPLY=()
  32. _get_comp_words_by_ref "$@" cur prev words cword
  33. }
  34. __%[1]s_index_of_word()
  35. {
  36. local w word=$1
  37. shift
  38. index=0
  39. for w in "$@"; do
  40. [[ $w = "$word" ]] && return
  41. index=$((index+1))
  42. done
  43. index=-1
  44. }
  45. __%[1]s_contains_word()
  46. {
  47. local w word=$1; shift
  48. for w in "$@"; do
  49. [[ $w = "$word" ]] && return
  50. done
  51. return 1
  52. }
  53. __%[1]s_handle_reply()
  54. {
  55. __%[1]s_debug "${FUNCNAME[0]}"
  56. case $cur in
  57. -*)
  58. if [[ $(type -t compopt) = "builtin" ]]; then
  59. compopt -o nospace
  60. fi
  61. local allflags
  62. if [ ${#must_have_one_flag[@]} -ne 0 ]; then
  63. allflags=("${must_have_one_flag[@]}")
  64. else
  65. allflags=("${flags[*]} ${two_word_flags[*]}")
  66. fi
  67. COMPREPLY=( $(compgen -W "${allflags[*]}" -- "$cur") )
  68. if [[ $(type -t compopt) = "builtin" ]]; then
  69. [[ "${COMPREPLY[0]}" == *= ]] || compopt +o nospace
  70. fi
  71. # complete after --flag=abc
  72. if [[ $cur == *=* ]]; then
  73. if [[ $(type -t compopt) = "builtin" ]]; then
  74. compopt +o nospace
  75. fi
  76. local index flag
  77. flag="${cur%%=*}"
  78. __%[1]s_index_of_word "${flag}" "${flags_with_completion[@]}"
  79. COMPREPLY=()
  80. if [[ ${index} -ge 0 ]]; then
  81. PREFIX=""
  82. cur="${cur#*=}"
  83. ${flags_completion[${index}]}
  84. if [ -n "${ZSH_VERSION}" ]; then
  85. # zsh completion needs --flag= prefix
  86. eval "COMPREPLY=( \"\${COMPREPLY[@]/#/${flag}=}\" )"
  87. fi
  88. fi
  89. fi
  90. return 0;
  91. ;;
  92. esac
  93. # check if we are handling a flag with special work handling
  94. local index
  95. __%[1]s_index_of_word "${prev}" "${flags_with_completion[@]}"
  96. if [[ ${index} -ge 0 ]]; then
  97. ${flags_completion[${index}]}
  98. return
  99. fi
  100. # we are parsing a flag and don't have a special handler, no completion
  101. if [[ ${cur} != "${words[cword]}" ]]; then
  102. return
  103. fi
  104. local completions
  105. completions=("${commands[@]}")
  106. if [[ ${#must_have_one_noun[@]} -ne 0 ]]; then
  107. completions=("${must_have_one_noun[@]}")
  108. fi
  109. if [[ ${#must_have_one_flag[@]} -ne 0 ]]; then
  110. completions+=("${must_have_one_flag[@]}")
  111. fi
  112. COMPREPLY=( $(compgen -W "${completions[*]}" -- "$cur") )
  113. if [[ ${#COMPREPLY[@]} -eq 0 && ${#noun_aliases[@]} -gt 0 && ${#must_have_one_noun[@]} -ne 0 ]]; then
  114. COMPREPLY=( $(compgen -W "${noun_aliases[*]}" -- "$cur") )
  115. fi
  116. if [[ ${#COMPREPLY[@]} -eq 0 ]]; then
  117. declare -F __custom_func >/dev/null && __custom_func
  118. fi
  119. # available in bash-completion >= 2, not always present on macOS
  120. if declare -F __ltrim_colon_completions >/dev/null; then
  121. __ltrim_colon_completions "$cur"
  122. fi
  123. # If there is only 1 completion and it is a flag with an = it will be completed
  124. # but we don't want a space after the =
  125. if [[ "${#COMPREPLY[@]}" -eq "1" ]] && [[ $(type -t compopt) = "builtin" ]] && [[ "${COMPREPLY[0]}" == --*= ]]; then
  126. compopt -o nospace
  127. fi
  128. }
  129. # The arguments should be in the form "ext1|ext2|extn"
  130. __%[1]s_handle_filename_extension_flag()
  131. {
  132. local ext="$1"
  133. _filedir "@(${ext})"
  134. }
  135. __%[1]s_handle_subdirs_in_dir_flag()
  136. {
  137. local dir="$1"
  138. pushd "${dir}" >/dev/null 2>&1 && _filedir -d && popd >/dev/null 2>&1
  139. }
  140. __%[1]s_handle_flag()
  141. {
  142. __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}"
  143. # if a command required a flag, and we found it, unset must_have_one_flag()
  144. local flagname=${words[c]}
  145. local flagvalue
  146. # if the word contained an =
  147. if [[ ${words[c]} == *"="* ]]; then
  148. flagvalue=${flagname#*=} # take in as flagvalue after the =
  149. flagname=${flagname%%=*} # strip everything after the =
  150. flagname="${flagname}=" # but put the = back
  151. fi
  152. __%[1]s_debug "${FUNCNAME[0]}: looking for ${flagname}"
  153. if __%[1]s_contains_word "${flagname}" "${must_have_one_flag[@]}"; then
  154. must_have_one_flag=()
  155. fi
  156. # if you set a flag which only applies to this command, don't show subcommands
  157. if __%[1]s_contains_word "${flagname}" "${local_nonpersistent_flags[@]}"; then
  158. commands=()
  159. fi
  160. # keep flag value with flagname as flaghash
  161. # flaghash variable is an associative array which is only supported in bash > 3.
  162. if [[ -z "${BASH_VERSION}" || "${BASH_VERSINFO[0]}" -gt 3 ]]; then
  163. if [ -n "${flagvalue}" ] ; then
  164. flaghash[${flagname}]=${flagvalue}
  165. elif [ -n "${words[ $((c+1)) ]}" ] ; then
  166. flaghash[${flagname}]=${words[ $((c+1)) ]}
  167. else
  168. flaghash[${flagname}]="true" # pad "true" for bool flag
  169. fi
  170. fi
  171. # skip the argument to a two word flag
  172. if __%[1]s_contains_word "${words[c]}" "${two_word_flags[@]}"; then
  173. c=$((c+1))
  174. # if we are looking for a flags value, don't show commands
  175. if [[ $c -eq $cword ]]; then
  176. commands=()
  177. fi
  178. fi
  179. c=$((c+1))
  180. }
  181. __%[1]s_handle_noun()
  182. {
  183. __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}"
  184. if __%[1]s_contains_word "${words[c]}" "${must_have_one_noun[@]}"; then
  185. must_have_one_noun=()
  186. elif __%[1]s_contains_word "${words[c]}" "${noun_aliases[@]}"; then
  187. must_have_one_noun=()
  188. fi
  189. nouns+=("${words[c]}")
  190. c=$((c+1))
  191. }
  192. __%[1]s_handle_command()
  193. {
  194. __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}"
  195. local next_command
  196. if [[ -n ${last_command} ]]; then
  197. next_command="_${last_command}_${words[c]//:/__}"
  198. else
  199. if [[ $c -eq 0 ]]; then
  200. next_command="_%[1]s_root_command"
  201. else
  202. next_command="_${words[c]//:/__}"
  203. fi
  204. fi
  205. c=$((c+1))
  206. __%[1]s_debug "${FUNCNAME[0]}: looking for ${next_command}"
  207. declare -F "$next_command" >/dev/null && $next_command
  208. }
  209. __%[1]s_handle_word()
  210. {
  211. if [[ $c -ge $cword ]]; then
  212. __%[1]s_handle_reply
  213. return
  214. fi
  215. __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}"
  216. if [[ "${words[c]}" == -* ]]; then
  217. __%[1]s_handle_flag
  218. elif __%[1]s_contains_word "${words[c]}" "${commands[@]}"; then
  219. __%[1]s_handle_command
  220. elif [[ $c -eq 0 ]]; then
  221. __%[1]s_handle_command
  222. elif __%[1]s_contains_word "${words[c]}" "${command_aliases[@]}"; then
  223. # aliashash variable is an associative array which is only supported in bash > 3.
  224. if [[ -z "${BASH_VERSION}" || "${BASH_VERSINFO[0]}" -gt 3 ]]; then
  225. words[c]=${aliashash[${words[c]}]}
  226. __%[1]s_handle_command
  227. else
  228. __%[1]s_handle_noun
  229. fi
  230. else
  231. __%[1]s_handle_noun
  232. fi
  233. __%[1]s_handle_word
  234. }
  235. `, name))
  236. }
  237. func writePostscript(buf *bytes.Buffer, name string) {
  238. name = strings.Replace(name, ":", "__", -1)
  239. buf.WriteString(fmt.Sprintf("__start_%s()\n", name))
  240. buf.WriteString(fmt.Sprintf(`{
  241. local cur prev words cword
  242. declare -A flaghash 2>/dev/null || :
  243. declare -A aliashash 2>/dev/null || :
  244. if declare -F _init_completion >/dev/null 2>&1; then
  245. _init_completion -s || return
  246. else
  247. __%[1]s_init_completion -n "=" || return
  248. fi
  249. local c=0
  250. local flags=()
  251. local two_word_flags=()
  252. local local_nonpersistent_flags=()
  253. local flags_with_completion=()
  254. local flags_completion=()
  255. local commands=("%[1]s")
  256. local must_have_one_flag=()
  257. local must_have_one_noun=()
  258. local last_command
  259. local nouns=()
  260. __%[1]s_handle_word
  261. }
  262. `, name))
  263. buf.WriteString(fmt.Sprintf(`if [[ $(type -t compopt) = "builtin" ]]; then
  264. complete -o default -F __start_%s %s
  265. else
  266. complete -o default -o nospace -F __start_%s %s
  267. fi
  268. `, name, name, name, name))
  269. buf.WriteString("# ex: ts=4 sw=4 et filetype=sh\n")
  270. }
  271. func writeCommands(buf *bytes.Buffer, cmd *Command) {
  272. buf.WriteString(" commands=()\n")
  273. for _, c := range cmd.Commands() {
  274. if !c.IsAvailableCommand() || c == cmd.helpCommand {
  275. continue
  276. }
  277. buf.WriteString(fmt.Sprintf(" commands+=(%q)\n", c.Name()))
  278. writeCmdAliases(buf, c)
  279. }
  280. buf.WriteString("\n")
  281. }
  282. func writeFlagHandler(buf *bytes.Buffer, name string, annotations map[string][]string, cmd *Command) {
  283. for key, value := range annotations {
  284. switch key {
  285. case BashCompFilenameExt:
  286. buf.WriteString(fmt.Sprintf(" flags_with_completion+=(%q)\n", name))
  287. var ext string
  288. if len(value) > 0 {
  289. ext = fmt.Sprintf("__%s_handle_filename_extension_flag ", cmd.Root().Name()) + strings.Join(value, "|")
  290. } else {
  291. ext = "_filedir"
  292. }
  293. buf.WriteString(fmt.Sprintf(" flags_completion+=(%q)\n", ext))
  294. case BashCompCustom:
  295. buf.WriteString(fmt.Sprintf(" flags_with_completion+=(%q)\n", name))
  296. if len(value) > 0 {
  297. handlers := strings.Join(value, "; ")
  298. buf.WriteString(fmt.Sprintf(" flags_completion+=(%q)\n", handlers))
  299. } else {
  300. buf.WriteString(" flags_completion+=(:)\n")
  301. }
  302. case BashCompSubdirsInDir:
  303. buf.WriteString(fmt.Sprintf(" flags_with_completion+=(%q)\n", name))
  304. var ext string
  305. if len(value) == 1 {
  306. ext = fmt.Sprintf("__%s_handle_subdirs_in_dir_flag ", cmd.Root().Name()) + value[0]
  307. } else {
  308. ext = "_filedir -d"
  309. }
  310. buf.WriteString(fmt.Sprintf(" flags_completion+=(%q)\n", ext))
  311. }
  312. }
  313. }
  314. func writeShortFlag(buf *bytes.Buffer, flag *pflag.Flag, cmd *Command) {
  315. name := flag.Shorthand
  316. format := " "
  317. if len(flag.NoOptDefVal) == 0 {
  318. format += "two_word_"
  319. }
  320. format += "flags+=(\"-%s\")\n"
  321. buf.WriteString(fmt.Sprintf(format, name))
  322. writeFlagHandler(buf, "-"+name, flag.Annotations, cmd)
  323. }
  324. func writeFlag(buf *bytes.Buffer, flag *pflag.Flag, cmd *Command) {
  325. name := flag.Name
  326. format := " flags+=(\"--%s"
  327. if len(flag.NoOptDefVal) == 0 {
  328. format += "="
  329. }
  330. format += "\")\n"
  331. buf.WriteString(fmt.Sprintf(format, name))
  332. writeFlagHandler(buf, "--"+name, flag.Annotations, cmd)
  333. }
  334. func writeLocalNonPersistentFlag(buf *bytes.Buffer, flag *pflag.Flag) {
  335. name := flag.Name
  336. format := " local_nonpersistent_flags+=(\"--%s"
  337. if len(flag.NoOptDefVal) == 0 {
  338. format += "="
  339. }
  340. format += "\")\n"
  341. buf.WriteString(fmt.Sprintf(format, name))
  342. }
  343. func writeFlags(buf *bytes.Buffer, cmd *Command) {
  344. buf.WriteString(` flags=()
  345. two_word_flags=()
  346. local_nonpersistent_flags=()
  347. flags_with_completion=()
  348. flags_completion=()
  349. `)
  350. localNonPersistentFlags := cmd.LocalNonPersistentFlags()
  351. cmd.NonInheritedFlags().VisitAll(func(flag *pflag.Flag) {
  352. if nonCompletableFlag(flag) {
  353. return
  354. }
  355. writeFlag(buf, flag, cmd)
  356. if len(flag.Shorthand) > 0 {
  357. writeShortFlag(buf, flag, cmd)
  358. }
  359. if localNonPersistentFlags.Lookup(flag.Name) != nil {
  360. writeLocalNonPersistentFlag(buf, flag)
  361. }
  362. })
  363. cmd.InheritedFlags().VisitAll(func(flag *pflag.Flag) {
  364. if nonCompletableFlag(flag) {
  365. return
  366. }
  367. writeFlag(buf, flag, cmd)
  368. if len(flag.Shorthand) > 0 {
  369. writeShortFlag(buf, flag, cmd)
  370. }
  371. })
  372. buf.WriteString("\n")
  373. }
  374. func writeRequiredFlag(buf *bytes.Buffer, cmd *Command) {
  375. buf.WriteString(" must_have_one_flag=()\n")
  376. flags := cmd.NonInheritedFlags()
  377. flags.VisitAll(func(flag *pflag.Flag) {
  378. if nonCompletableFlag(flag) {
  379. return
  380. }
  381. for key := range flag.Annotations {
  382. switch key {
  383. case BashCompOneRequiredFlag:
  384. format := " must_have_one_flag+=(\"--%s"
  385. if flag.Value.Type() != "bool" {
  386. format += "="
  387. }
  388. format += "\")\n"
  389. buf.WriteString(fmt.Sprintf(format, flag.Name))
  390. if len(flag.Shorthand) > 0 {
  391. buf.WriteString(fmt.Sprintf(" must_have_one_flag+=(\"-%s\")\n", flag.Shorthand))
  392. }
  393. }
  394. }
  395. })
  396. }
  397. func writeRequiredNouns(buf *bytes.Buffer, cmd *Command) {
  398. buf.WriteString(" must_have_one_noun=()\n")
  399. sort.Sort(sort.StringSlice(cmd.ValidArgs))
  400. for _, value := range cmd.ValidArgs {
  401. buf.WriteString(fmt.Sprintf(" must_have_one_noun+=(%q)\n", value))
  402. }
  403. }
  404. func writeCmdAliases(buf *bytes.Buffer, cmd *Command) {
  405. if len(cmd.Aliases) == 0 {
  406. return
  407. }
  408. sort.Sort(sort.StringSlice(cmd.Aliases))
  409. buf.WriteString(fmt.Sprint(` if [[ -z "${BASH_VERSION}" || "${BASH_VERSINFO[0]}" -gt 3 ]]; then`, "\n"))
  410. for _, value := range cmd.Aliases {
  411. buf.WriteString(fmt.Sprintf(" command_aliases+=(%q)\n", value))
  412. buf.WriteString(fmt.Sprintf(" aliashash[%q]=%q\n", value, cmd.Name()))
  413. }
  414. buf.WriteString(` fi`)
  415. buf.WriteString("\n")
  416. }
  417. func writeArgAliases(buf *bytes.Buffer, cmd *Command) {
  418. buf.WriteString(" noun_aliases=()\n")
  419. sort.Sort(sort.StringSlice(cmd.ArgAliases))
  420. for _, value := range cmd.ArgAliases {
  421. buf.WriteString(fmt.Sprintf(" noun_aliases+=(%q)\n", value))
  422. }
  423. }
  424. func gen(buf *bytes.Buffer, cmd *Command) {
  425. for _, c := range cmd.Commands() {
  426. if !c.IsAvailableCommand() || c == cmd.helpCommand {
  427. continue
  428. }
  429. gen(buf, c)
  430. }
  431. commandName := cmd.CommandPath()
  432. commandName = strings.Replace(commandName, " ", "_", -1)
  433. commandName = strings.Replace(commandName, ":", "__", -1)
  434. if cmd.Root() == cmd {
  435. buf.WriteString(fmt.Sprintf("_%s_root_command()\n{\n", commandName))
  436. } else {
  437. buf.WriteString(fmt.Sprintf("_%s()\n{\n", commandName))
  438. }
  439. buf.WriteString(fmt.Sprintf(" last_command=%q\n", commandName))
  440. buf.WriteString("\n")
  441. buf.WriteString(" command_aliases=()\n")
  442. buf.WriteString("\n")
  443. writeCommands(buf, cmd)
  444. writeFlags(buf, cmd)
  445. writeRequiredFlag(buf, cmd)
  446. writeRequiredNouns(buf, cmd)
  447. writeArgAliases(buf, cmd)
  448. buf.WriteString("}\n\n")
  449. }
  450. // GenBashCompletion generates bash completion file and writes to the passed writer.
  451. func (c *Command) GenBashCompletion(w io.Writer) error {
  452. buf := new(bytes.Buffer)
  453. writePreamble(buf, c.Name())
  454. if len(c.BashCompletionFunction) > 0 {
  455. buf.WriteString(c.BashCompletionFunction + "\n")
  456. }
  457. gen(buf, c)
  458. writePostscript(buf, c.Name())
  459. _, err := buf.WriteTo(w)
  460. return err
  461. }
  462. func nonCompletableFlag(flag *pflag.Flag) bool {
  463. return flag.Hidden || len(flag.Deprecated) > 0
  464. }
  465. // GenBashCompletionFile generates bash completion file.
  466. func (c *Command) GenBashCompletionFile(filename string) error {
  467. outFile, err := os.Create(filename)
  468. if err != nil {
  469. return err
  470. }
  471. defer outFile.Close()
  472. return c.GenBashCompletion(outFile)
  473. }
  474. // MarkFlagRequired adds the BashCompOneRequiredFlag annotation to the named flag if it exists,
  475. // and causes your command to report an error if invoked without the flag.
  476. func (c *Command) MarkFlagRequired(name string) error {
  477. return MarkFlagRequired(c.Flags(), name)
  478. }
  479. // MarkPersistentFlagRequired adds the BashCompOneRequiredFlag annotation to the named persistent flag if it exists,
  480. // and causes your command to report an error if invoked without the flag.
  481. func (c *Command) MarkPersistentFlagRequired(name string) error {
  482. return MarkFlagRequired(c.PersistentFlags(), name)
  483. }
  484. // MarkFlagRequired adds the BashCompOneRequiredFlag annotation to the named flag if it exists,
  485. // and causes your command to report an error if invoked without the flag.
  486. func MarkFlagRequired(flags *pflag.FlagSet, name string) error {
  487. return flags.SetAnnotation(name, BashCompOneRequiredFlag, []string{"true"})
  488. }
  489. // MarkFlagFilename adds the BashCompFilenameExt annotation to the named flag, if it exists.
  490. // Generated bash autocompletion will select filenames for the flag, limiting to named extensions if provided.
  491. func (c *Command) MarkFlagFilename(name string, extensions ...string) error {
  492. return MarkFlagFilename(c.Flags(), name, extensions...)
  493. }
  494. // MarkFlagCustom adds the BashCompCustom annotation to the named flag, if it exists.
  495. // Generated bash autocompletion will call the bash function f for the flag.
  496. func (c *Command) MarkFlagCustom(name string, f string) error {
  497. return MarkFlagCustom(c.Flags(), name, f)
  498. }
  499. // MarkPersistentFlagFilename adds the BashCompFilenameExt annotation to the named persistent flag, if it exists.
  500. // Generated bash autocompletion will select filenames for the flag, limiting to named extensions if provided.
  501. func (c *Command) MarkPersistentFlagFilename(name string, extensions ...string) error {
  502. return MarkFlagFilename(c.PersistentFlags(), name, extensions...)
  503. }
  504. // MarkFlagFilename adds the BashCompFilenameExt annotation to the named flag in the flag set, if it exists.
  505. // Generated bash autocompletion will select filenames for the flag, limiting to named extensions if provided.
  506. func MarkFlagFilename(flags *pflag.FlagSet, name string, extensions ...string) error {
  507. return flags.SetAnnotation(name, BashCompFilenameExt, extensions)
  508. }
  509. // MarkFlagCustom adds the BashCompCustom annotation to the named flag in the flag set, if it exists.
  510. // Generated bash autocompletion will call the bash function f for the flag.
  511. func MarkFlagCustom(flags *pflag.FlagSet, name string, f string) error {
  512. return flags.SetAnnotation(name, BashCompCustom, []string{f})
  513. }