bash_completions.go 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548
  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. if declare -F __%[1]s_custom_func >/dev/null; then
  118. # try command name qualified custom func
  119. __%[1]s_custom_func
  120. else
  121. # otherwise fall back to unqualified for compatibility
  122. declare -F __custom_func >/dev/null && __custom_func
  123. fi
  124. fi
  125. # available in bash-completion >= 2, not always present on macOS
  126. if declare -F __ltrim_colon_completions >/dev/null; then
  127. __ltrim_colon_completions "$cur"
  128. fi
  129. # If there is only 1 completion and it is a flag with an = it will be completed
  130. # but we don't want a space after the =
  131. if [[ "${#COMPREPLY[@]}" -eq "1" ]] && [[ $(type -t compopt) = "builtin" ]] && [[ "${COMPREPLY[0]}" == --*= ]]; then
  132. compopt -o nospace
  133. fi
  134. }
  135. # The arguments should be in the form "ext1|ext2|extn"
  136. __%[1]s_handle_filename_extension_flag()
  137. {
  138. local ext="$1"
  139. _filedir "@(${ext})"
  140. }
  141. __%[1]s_handle_subdirs_in_dir_flag()
  142. {
  143. local dir="$1"
  144. pushd "${dir}" >/dev/null 2>&1 && _filedir -d && popd >/dev/null 2>&1
  145. }
  146. __%[1]s_handle_flag()
  147. {
  148. __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}"
  149. # if a command required a flag, and we found it, unset must_have_one_flag()
  150. local flagname=${words[c]}
  151. local flagvalue
  152. # if the word contained an =
  153. if [[ ${words[c]} == *"="* ]]; then
  154. flagvalue=${flagname#*=} # take in as flagvalue after the =
  155. flagname=${flagname%%=*} # strip everything after the =
  156. flagname="${flagname}=" # but put the = back
  157. fi
  158. __%[1]s_debug "${FUNCNAME[0]}: looking for ${flagname}"
  159. if __%[1]s_contains_word "${flagname}" "${must_have_one_flag[@]}"; then
  160. must_have_one_flag=()
  161. fi
  162. # if you set a flag which only applies to this command, don't show subcommands
  163. if __%[1]s_contains_word "${flagname}" "${local_nonpersistent_flags[@]}"; then
  164. commands=()
  165. fi
  166. # keep flag value with flagname as flaghash
  167. # flaghash variable is an associative array which is only supported in bash > 3.
  168. if [[ -z "${BASH_VERSION}" || "${BASH_VERSINFO[0]}" -gt 3 ]]; then
  169. if [ -n "${flagvalue}" ] ; then
  170. flaghash[${flagname}]=${flagvalue}
  171. elif [ -n "${words[ $((c+1)) ]}" ] ; then
  172. flaghash[${flagname}]=${words[ $((c+1)) ]}
  173. else
  174. flaghash[${flagname}]="true" # pad "true" for bool flag
  175. fi
  176. fi
  177. # skip the argument to a two word flag
  178. if [[ ${words[c]} != *"="* ]] && __%[1]s_contains_word "${words[c]}" "${two_word_flags[@]}"; then
  179. __%[1]s_debug "${FUNCNAME[0]}: found a flag ${words[c]}, skip the next argument"
  180. c=$((c+1))
  181. # if we are looking for a flags value, don't show commands
  182. if [[ $c -eq $cword ]]; then
  183. commands=()
  184. fi
  185. fi
  186. c=$((c+1))
  187. }
  188. __%[1]s_handle_noun()
  189. {
  190. __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}"
  191. if __%[1]s_contains_word "${words[c]}" "${must_have_one_noun[@]}"; then
  192. must_have_one_noun=()
  193. elif __%[1]s_contains_word "${words[c]}" "${noun_aliases[@]}"; then
  194. must_have_one_noun=()
  195. fi
  196. nouns+=("${words[c]}")
  197. c=$((c+1))
  198. }
  199. __%[1]s_handle_command()
  200. {
  201. __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}"
  202. local next_command
  203. if [[ -n ${last_command} ]]; then
  204. next_command="_${last_command}_${words[c]//:/__}"
  205. else
  206. if [[ $c -eq 0 ]]; then
  207. next_command="_%[1]s_root_command"
  208. else
  209. next_command="_${words[c]//:/__}"
  210. fi
  211. fi
  212. c=$((c+1))
  213. __%[1]s_debug "${FUNCNAME[0]}: looking for ${next_command}"
  214. declare -F "$next_command" >/dev/null && $next_command
  215. }
  216. __%[1]s_handle_word()
  217. {
  218. if [[ $c -ge $cword ]]; then
  219. __%[1]s_handle_reply
  220. return
  221. fi
  222. __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}"
  223. if [[ "${words[c]}" == -* ]]; then
  224. __%[1]s_handle_flag
  225. elif __%[1]s_contains_word "${words[c]}" "${commands[@]}"; then
  226. __%[1]s_handle_command
  227. elif [[ $c -eq 0 ]]; then
  228. __%[1]s_handle_command
  229. elif __%[1]s_contains_word "${words[c]}" "${command_aliases[@]}"; then
  230. # aliashash variable is an associative array which is only supported in bash > 3.
  231. if [[ -z "${BASH_VERSION}" || "${BASH_VERSINFO[0]}" -gt 3 ]]; then
  232. words[c]=${aliashash[${words[c]}]}
  233. __%[1]s_handle_command
  234. else
  235. __%[1]s_handle_noun
  236. fi
  237. else
  238. __%[1]s_handle_noun
  239. fi
  240. __%[1]s_handle_word
  241. }
  242. `, name))
  243. }
  244. func writePostscript(buf *bytes.Buffer, name string) {
  245. name = strings.Replace(name, ":", "__", -1)
  246. buf.WriteString(fmt.Sprintf("__start_%s()\n", name))
  247. buf.WriteString(fmt.Sprintf(`{
  248. local cur prev words cword
  249. declare -A flaghash 2>/dev/null || :
  250. declare -A aliashash 2>/dev/null || :
  251. if declare -F _init_completion >/dev/null 2>&1; then
  252. _init_completion -s || return
  253. else
  254. __%[1]s_init_completion -n "=" || return
  255. fi
  256. local c=0
  257. local flags=()
  258. local two_word_flags=()
  259. local local_nonpersistent_flags=()
  260. local flags_with_completion=()
  261. local flags_completion=()
  262. local commands=("%[1]s")
  263. local must_have_one_flag=()
  264. local must_have_one_noun=()
  265. local last_command
  266. local nouns=()
  267. __%[1]s_handle_word
  268. }
  269. `, name))
  270. buf.WriteString(fmt.Sprintf(`if [[ $(type -t compopt) = "builtin" ]]; then
  271. complete -o default -F __start_%s %s
  272. else
  273. complete -o default -o nospace -F __start_%s %s
  274. fi
  275. `, name, name, name, name))
  276. buf.WriteString("# ex: ts=4 sw=4 et filetype=sh\n")
  277. }
  278. func writeCommands(buf *bytes.Buffer, cmd *Command) {
  279. buf.WriteString(" commands=()\n")
  280. for _, c := range cmd.Commands() {
  281. if !c.IsAvailableCommand() || c == cmd.helpCommand {
  282. continue
  283. }
  284. buf.WriteString(fmt.Sprintf(" commands+=(%q)\n", c.Name()))
  285. writeCmdAliases(buf, c)
  286. }
  287. buf.WriteString("\n")
  288. }
  289. func writeFlagHandler(buf *bytes.Buffer, name string, annotations map[string][]string, cmd *Command) {
  290. for key, value := range annotations {
  291. switch key {
  292. case BashCompFilenameExt:
  293. buf.WriteString(fmt.Sprintf(" flags_with_completion+=(%q)\n", name))
  294. var ext string
  295. if len(value) > 0 {
  296. ext = fmt.Sprintf("__%s_handle_filename_extension_flag ", cmd.Root().Name()) + strings.Join(value, "|")
  297. } else {
  298. ext = "_filedir"
  299. }
  300. buf.WriteString(fmt.Sprintf(" flags_completion+=(%q)\n", ext))
  301. case BashCompCustom:
  302. buf.WriteString(fmt.Sprintf(" flags_with_completion+=(%q)\n", name))
  303. if len(value) > 0 {
  304. handlers := strings.Join(value, "; ")
  305. buf.WriteString(fmt.Sprintf(" flags_completion+=(%q)\n", handlers))
  306. } else {
  307. buf.WriteString(" flags_completion+=(:)\n")
  308. }
  309. case BashCompSubdirsInDir:
  310. buf.WriteString(fmt.Sprintf(" flags_with_completion+=(%q)\n", name))
  311. var ext string
  312. if len(value) == 1 {
  313. ext = fmt.Sprintf("__%s_handle_subdirs_in_dir_flag ", cmd.Root().Name()) + value[0]
  314. } else {
  315. ext = "_filedir -d"
  316. }
  317. buf.WriteString(fmt.Sprintf(" flags_completion+=(%q)\n", ext))
  318. }
  319. }
  320. }
  321. func writeShortFlag(buf *bytes.Buffer, flag *pflag.Flag, cmd *Command) {
  322. name := flag.Shorthand
  323. format := " "
  324. if len(flag.NoOptDefVal) == 0 {
  325. format += "two_word_"
  326. }
  327. format += "flags+=(\"-%s\")\n"
  328. buf.WriteString(fmt.Sprintf(format, name))
  329. writeFlagHandler(buf, "-"+name, flag.Annotations, cmd)
  330. }
  331. func writeFlag(buf *bytes.Buffer, flag *pflag.Flag, cmd *Command) {
  332. name := flag.Name
  333. format := " flags+=(\"--%s"
  334. if len(flag.NoOptDefVal) == 0 {
  335. format += "="
  336. }
  337. format += "\")\n"
  338. buf.WriteString(fmt.Sprintf(format, name))
  339. if len(flag.NoOptDefVal) == 0 {
  340. format = " two_word_flags+=(\"--%s\")\n"
  341. buf.WriteString(fmt.Sprintf(format, name))
  342. }
  343. writeFlagHandler(buf, "--"+name, flag.Annotations, cmd)
  344. }
  345. func writeLocalNonPersistentFlag(buf *bytes.Buffer, flag *pflag.Flag) {
  346. name := flag.Name
  347. format := " local_nonpersistent_flags+=(\"--%s"
  348. if len(flag.NoOptDefVal) == 0 {
  349. format += "="
  350. }
  351. format += "\")\n"
  352. buf.WriteString(fmt.Sprintf(format, name))
  353. }
  354. func writeFlags(buf *bytes.Buffer, cmd *Command) {
  355. buf.WriteString(` flags=()
  356. two_word_flags=()
  357. local_nonpersistent_flags=()
  358. flags_with_completion=()
  359. flags_completion=()
  360. `)
  361. localNonPersistentFlags := cmd.LocalNonPersistentFlags()
  362. cmd.NonInheritedFlags().VisitAll(func(flag *pflag.Flag) {
  363. if nonCompletableFlag(flag) {
  364. return
  365. }
  366. writeFlag(buf, flag, cmd)
  367. if len(flag.Shorthand) > 0 {
  368. writeShortFlag(buf, flag, cmd)
  369. }
  370. if localNonPersistentFlags.Lookup(flag.Name) != nil {
  371. writeLocalNonPersistentFlag(buf, flag)
  372. }
  373. })
  374. cmd.InheritedFlags().VisitAll(func(flag *pflag.Flag) {
  375. if nonCompletableFlag(flag) {
  376. return
  377. }
  378. writeFlag(buf, flag, cmd)
  379. if len(flag.Shorthand) > 0 {
  380. writeShortFlag(buf, flag, cmd)
  381. }
  382. })
  383. buf.WriteString("\n")
  384. }
  385. func writeRequiredFlag(buf *bytes.Buffer, cmd *Command) {
  386. buf.WriteString(" must_have_one_flag=()\n")
  387. flags := cmd.NonInheritedFlags()
  388. flags.VisitAll(func(flag *pflag.Flag) {
  389. if nonCompletableFlag(flag) {
  390. return
  391. }
  392. for key := range flag.Annotations {
  393. switch key {
  394. case BashCompOneRequiredFlag:
  395. format := " must_have_one_flag+=(\"--%s"
  396. if flag.Value.Type() != "bool" {
  397. format += "="
  398. }
  399. format += "\")\n"
  400. buf.WriteString(fmt.Sprintf(format, flag.Name))
  401. if len(flag.Shorthand) > 0 {
  402. buf.WriteString(fmt.Sprintf(" must_have_one_flag+=(\"-%s\")\n", flag.Shorthand))
  403. }
  404. }
  405. }
  406. })
  407. }
  408. func writeRequiredNouns(buf *bytes.Buffer, cmd *Command) {
  409. buf.WriteString(" must_have_one_noun=()\n")
  410. sort.Sort(sort.StringSlice(cmd.ValidArgs))
  411. for _, value := range cmd.ValidArgs {
  412. buf.WriteString(fmt.Sprintf(" must_have_one_noun+=(%q)\n", value))
  413. }
  414. }
  415. func writeCmdAliases(buf *bytes.Buffer, cmd *Command) {
  416. if len(cmd.Aliases) == 0 {
  417. return
  418. }
  419. sort.Sort(sort.StringSlice(cmd.Aliases))
  420. buf.WriteString(fmt.Sprint(` if [[ -z "${BASH_VERSION}" || "${BASH_VERSINFO[0]}" -gt 3 ]]; then`, "\n"))
  421. for _, value := range cmd.Aliases {
  422. buf.WriteString(fmt.Sprintf(" command_aliases+=(%q)\n", value))
  423. buf.WriteString(fmt.Sprintf(" aliashash[%q]=%q\n", value, cmd.Name()))
  424. }
  425. buf.WriteString(` fi`)
  426. buf.WriteString("\n")
  427. }
  428. func writeArgAliases(buf *bytes.Buffer, cmd *Command) {
  429. buf.WriteString(" noun_aliases=()\n")
  430. sort.Sort(sort.StringSlice(cmd.ArgAliases))
  431. for _, value := range cmd.ArgAliases {
  432. buf.WriteString(fmt.Sprintf(" noun_aliases+=(%q)\n", value))
  433. }
  434. }
  435. func gen(buf *bytes.Buffer, cmd *Command) {
  436. for _, c := range cmd.Commands() {
  437. if !c.IsAvailableCommand() || c == cmd.helpCommand {
  438. continue
  439. }
  440. gen(buf, c)
  441. }
  442. commandName := cmd.CommandPath()
  443. commandName = strings.Replace(commandName, " ", "_", -1)
  444. commandName = strings.Replace(commandName, ":", "__", -1)
  445. if cmd.Root() == cmd {
  446. buf.WriteString(fmt.Sprintf("_%s_root_command()\n{\n", commandName))
  447. } else {
  448. buf.WriteString(fmt.Sprintf("_%s()\n{\n", commandName))
  449. }
  450. buf.WriteString(fmt.Sprintf(" last_command=%q\n", commandName))
  451. buf.WriteString("\n")
  452. buf.WriteString(" command_aliases=()\n")
  453. buf.WriteString("\n")
  454. writeCommands(buf, cmd)
  455. writeFlags(buf, cmd)
  456. writeRequiredFlag(buf, cmd)
  457. writeRequiredNouns(buf, cmd)
  458. writeArgAliases(buf, cmd)
  459. buf.WriteString("}\n\n")
  460. }
  461. // GenBashCompletion generates bash completion file and writes to the passed writer.
  462. func (c *Command) GenBashCompletion(w io.Writer) error {
  463. buf := new(bytes.Buffer)
  464. writePreamble(buf, c.Name())
  465. if len(c.BashCompletionFunction) > 0 {
  466. buf.WriteString(c.BashCompletionFunction + "\n")
  467. }
  468. gen(buf, c)
  469. writePostscript(buf, c.Name())
  470. _, err := buf.WriteTo(w)
  471. return err
  472. }
  473. func nonCompletableFlag(flag *pflag.Flag) bool {
  474. return flag.Hidden || len(flag.Deprecated) > 0
  475. }
  476. // GenBashCompletionFile generates bash completion file.
  477. func (c *Command) GenBashCompletionFile(filename string) error {
  478. outFile, err := os.Create(filename)
  479. if err != nil {
  480. return err
  481. }
  482. defer outFile.Close()
  483. return c.GenBashCompletion(outFile)
  484. }