1package cobra 2 3import ( 4 "bytes" 5 "fmt" 6 "io" 7 "os" 8 "sort" 9 "strings" 10 11 "github.com/spf13/pflag" 12) 13 14// Annotations for Bash completion. 15const ( 16 BashCompFilenameExt = "cobra_annotation_bash_completion_filename_extensions" 17 BashCompCustom = "cobra_annotation_bash_completion_custom" 18 BashCompOneRequiredFlag = "cobra_annotation_bash_completion_one_required_flag" 19 BashCompSubdirsInDir = "cobra_annotation_bash_completion_subdirs_in_dir" 20) 21 22func writePreamble(buf *bytes.Buffer, name string) { 23 buf.WriteString(fmt.Sprintf("# bash completion for %-36s -*- shell-script -*-\n", name)) 24 buf.WriteString(fmt.Sprintf(` 25__%[1]s_debug() 26{ 27 if [[ -n ${BASH_COMP_DEBUG_FILE} ]]; then 28 echo "$*" >> "${BASH_COMP_DEBUG_FILE}" 29 fi 30} 31 32# Homebrew on Macs have version 1.3 of bash-completion which doesn't include 33# _init_completion. This is a very minimal version of that function. 34__%[1]s_init_completion() 35{ 36 COMPREPLY=() 37 _get_comp_words_by_ref "$@" cur prev words cword 38} 39 40__%[1]s_index_of_word() 41{ 42 local w word=$1 43 shift 44 index=0 45 for w in "$@"; do 46 [[ $w = "$word" ]] && return 47 index=$((index+1)) 48 done 49 index=-1 50} 51 52__%[1]s_contains_word() 53{ 54 local w word=$1; shift 55 for w in "$@"; do 56 [[ $w = "$word" ]] && return 57 done 58 return 1 59} 60 61__%[1]s_handle_go_custom_completion() 62{ 63 __%[1]s_debug "${FUNCNAME[0]}: cur is ${cur}, words[*] is ${words[*]}, #words[@] is ${#words[@]}" 64 65 local out requestComp lastParam lastChar comp directive args 66 67 # Prepare the command to request completions for the program. 68 # Calling ${words[0]} instead of directly %[1]s allows to handle aliases 69 args=("${words[@]:1}") 70 requestComp="${words[0]} %[2]s ${args[*]}" 71 72 lastParam=${words[$((${#words[@]}-1))]} 73 lastChar=${lastParam:$((${#lastParam}-1)):1} 74 __%[1]s_debug "${FUNCNAME[0]}: lastParam ${lastParam}, lastChar ${lastChar}" 75 76 if [ -z "${cur}" ] && [ "${lastChar}" != "=" ]; then 77 # If the last parameter is complete (there is a space following it) 78 # We add an extra empty parameter so we can indicate this to the go method. 79 __%[1]s_debug "${FUNCNAME[0]}: Adding extra empty parameter" 80 requestComp="${requestComp} \"\"" 81 fi 82 83 __%[1]s_debug "${FUNCNAME[0]}: calling ${requestComp}" 84 # Use eval to handle any environment variables and such 85 out=$(eval "${requestComp}" 2>/dev/null) 86 87 # Extract the directive integer at the very end of the output following a colon (:) 88 directive=${out##*:} 89 # Remove the directive 90 out=${out%%:*} 91 if [ "${directive}" = "${out}" ]; then 92 # There is not directive specified 93 directive=0 94 fi 95 __%[1]s_debug "${FUNCNAME[0]}: the completion directive is: ${directive}" 96 __%[1]s_debug "${FUNCNAME[0]}: the completions are: ${out[*]}" 97 98 if [ $((directive & %[3]d)) -ne 0 ]; then 99 # Error code. No completion. 100 __%[1]s_debug "${FUNCNAME[0]}: received error from custom completion go code" 101 return 102 else 103 if [ $((directive & %[4]d)) -ne 0 ]; then 104 if [[ $(type -t compopt) = "builtin" ]]; then 105 __%[1]s_debug "${FUNCNAME[0]}: activating no space" 106 compopt -o nospace 107 fi 108 fi 109 if [ $((directive & %[5]d)) -ne 0 ]; then 110 if [[ $(type -t compopt) = "builtin" ]]; then 111 __%[1]s_debug "${FUNCNAME[0]}: activating no file completion" 112 compopt +o default 113 fi 114 fi 115 116 while IFS='' read -r comp; do 117 COMPREPLY+=("$comp") 118 done < <(compgen -W "${out[*]}" -- "$cur") 119 fi 120} 121 122__%[1]s_handle_reply() 123{ 124 __%[1]s_debug "${FUNCNAME[0]}" 125 local comp 126 case $cur in 127 -*) 128 if [[ $(type -t compopt) = "builtin" ]]; then 129 compopt -o nospace 130 fi 131 local allflags 132 if [ ${#must_have_one_flag[@]} -ne 0 ]; then 133 allflags=("${must_have_one_flag[@]}") 134 else 135 allflags=("${flags[*]} ${two_word_flags[*]}") 136 fi 137 while IFS='' read -r comp; do 138 COMPREPLY+=("$comp") 139 done < <(compgen -W "${allflags[*]}" -- "$cur") 140 if [[ $(type -t compopt) = "builtin" ]]; then 141 [[ "${COMPREPLY[0]}" == *= ]] || compopt +o nospace 142 fi 143 144 # complete after --flag=abc 145 if [[ $cur == *=* ]]; then 146 if [[ $(type -t compopt) = "builtin" ]]; then 147 compopt +o nospace 148 fi 149 150 local index flag 151 flag="${cur%%=*}" 152 __%[1]s_index_of_word "${flag}" "${flags_with_completion[@]}" 153 COMPREPLY=() 154 if [[ ${index} -ge 0 ]]; then 155 PREFIX="" 156 cur="${cur#*=}" 157 ${flags_completion[${index}]} 158 if [ -n "${ZSH_VERSION}" ]; then 159 # zsh completion needs --flag= prefix 160 eval "COMPREPLY=( \"\${COMPREPLY[@]/#/${flag}=}\" )" 161 fi 162 fi 163 fi 164 return 0; 165 ;; 166 esac 167 168 # check if we are handling a flag with special work handling 169 local index 170 __%[1]s_index_of_word "${prev}" "${flags_with_completion[@]}" 171 if [[ ${index} -ge 0 ]]; then 172 ${flags_completion[${index}]} 173 return 174 fi 175 176 # we are parsing a flag and don't have a special handler, no completion 177 if [[ ${cur} != "${words[cword]}" ]]; then 178 return 179 fi 180 181 local completions 182 completions=("${commands[@]}") 183 if [[ ${#must_have_one_noun[@]} -ne 0 ]]; then 184 completions=("${must_have_one_noun[@]}") 185 elif [[ -n "${has_completion_function}" ]]; then 186 # if a go completion function is provided, defer to that function 187 completions=() 188 __%[1]s_handle_go_custom_completion 189 fi 190 if [[ ${#must_have_one_flag[@]} -ne 0 ]]; then 191 completions+=("${must_have_one_flag[@]}") 192 fi 193 while IFS='' read -r comp; do 194 COMPREPLY+=("$comp") 195 done < <(compgen -W "${completions[*]}" -- "$cur") 196 197 if [[ ${#COMPREPLY[@]} -eq 0 && ${#noun_aliases[@]} -gt 0 && ${#must_have_one_noun[@]} -ne 0 ]]; then 198 while IFS='' read -r comp; do 199 COMPREPLY+=("$comp") 200 done < <(compgen -W "${noun_aliases[*]}" -- "$cur") 201 fi 202 203 if [[ ${#COMPREPLY[@]} -eq 0 ]]; then 204 if declare -F __%[1]s_custom_func >/dev/null; then 205 # try command name qualified custom func 206 __%[1]s_custom_func 207 else 208 # otherwise fall back to unqualified for compatibility 209 declare -F __custom_func >/dev/null && __custom_func 210 fi 211 fi 212 213 # available in bash-completion >= 2, not always present on macOS 214 if declare -F __ltrim_colon_completions >/dev/null; then 215 __ltrim_colon_completions "$cur" 216 fi 217 218 # If there is only 1 completion and it is a flag with an = it will be completed 219 # but we don't want a space after the = 220 if [[ "${#COMPREPLY[@]}" -eq "1" ]] && [[ $(type -t compopt) = "builtin" ]] && [[ "${COMPREPLY[0]}" == --*= ]]; then 221 compopt -o nospace 222 fi 223} 224 225# The arguments should be in the form "ext1|ext2|extn" 226__%[1]s_handle_filename_extension_flag() 227{ 228 local ext="$1" 229 _filedir "@(${ext})" 230} 231 232__%[1]s_handle_subdirs_in_dir_flag() 233{ 234 local dir="$1" 235 pushd "${dir}" >/dev/null 2>&1 && _filedir -d && popd >/dev/null 2>&1 || return 236} 237 238__%[1]s_handle_flag() 239{ 240 __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}" 241 242 # if a command required a flag, and we found it, unset must_have_one_flag() 243 local flagname=${words[c]} 244 local flagvalue 245 # if the word contained an = 246 if [[ ${words[c]} == *"="* ]]; then 247 flagvalue=${flagname#*=} # take in as flagvalue after the = 248 flagname=${flagname%%=*} # strip everything after the = 249 flagname="${flagname}=" # but put the = back 250 fi 251 __%[1]s_debug "${FUNCNAME[0]}: looking for ${flagname}" 252 if __%[1]s_contains_word "${flagname}" "${must_have_one_flag[@]}"; then 253 must_have_one_flag=() 254 fi 255 256 # if you set a flag which only applies to this command, don't show subcommands 257 if __%[1]s_contains_word "${flagname}" "${local_nonpersistent_flags[@]}"; then 258 commands=() 259 fi 260 261 # keep flag value with flagname as flaghash 262 # flaghash variable is an associative array which is only supported in bash > 3. 263 if [[ -z "${BASH_VERSION}" || "${BASH_VERSINFO[0]}" -gt 3 ]]; then 264 if [ -n "${flagvalue}" ] ; then 265 flaghash[${flagname}]=${flagvalue} 266 elif [ -n "${words[ $((c+1)) ]}" ] ; then 267 flaghash[${flagname}]=${words[ $((c+1)) ]} 268 else 269 flaghash[${flagname}]="true" # pad "true" for bool flag 270 fi 271 fi 272 273 # skip the argument to a two word flag 274 if [[ ${words[c]} != *"="* ]] && __%[1]s_contains_word "${words[c]}" "${two_word_flags[@]}"; then 275 __%[1]s_debug "${FUNCNAME[0]}: found a flag ${words[c]}, skip the next argument" 276 c=$((c+1)) 277 # if we are looking for a flags value, don't show commands 278 if [[ $c -eq $cword ]]; then 279 commands=() 280 fi 281 fi 282 283 c=$((c+1)) 284 285} 286 287__%[1]s_handle_noun() 288{ 289 __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}" 290 291 if __%[1]s_contains_word "${words[c]}" "${must_have_one_noun[@]}"; then 292 must_have_one_noun=() 293 elif __%[1]s_contains_word "${words[c]}" "${noun_aliases[@]}"; then 294 must_have_one_noun=() 295 fi 296 297 nouns+=("${words[c]}") 298 c=$((c+1)) 299} 300 301__%[1]s_handle_command() 302{ 303 __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}" 304 305 local next_command 306 if [[ -n ${last_command} ]]; then 307 next_command="_${last_command}_${words[c]//:/__}" 308 else 309 if [[ $c -eq 0 ]]; then 310 next_command="_%[1]s_root_command" 311 else 312 next_command="_${words[c]//:/__}" 313 fi 314 fi 315 c=$((c+1)) 316 __%[1]s_debug "${FUNCNAME[0]}: looking for ${next_command}" 317 declare -F "$next_command" >/dev/null && $next_command 318} 319 320__%[1]s_handle_word() 321{ 322 if [[ $c -ge $cword ]]; then 323 __%[1]s_handle_reply 324 return 325 fi 326 __%[1]s_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}" 327 if [[ "${words[c]}" == -* ]]; then 328 __%[1]s_handle_flag 329 elif __%[1]s_contains_word "${words[c]}" "${commands[@]}"; then 330 __%[1]s_handle_command 331 elif [[ $c -eq 0 ]]; then 332 __%[1]s_handle_command 333 elif __%[1]s_contains_word "${words[c]}" "${command_aliases[@]}"; then 334 # aliashash variable is an associative array which is only supported in bash > 3. 335 if [[ -z "${BASH_VERSION}" || "${BASH_VERSINFO[0]}" -gt 3 ]]; then 336 words[c]=${aliashash[${words[c]}]} 337 __%[1]s_handle_command 338 else 339 __%[1]s_handle_noun 340 fi 341 else 342 __%[1]s_handle_noun 343 fi 344 __%[1]s_handle_word 345} 346 347`, name, ShellCompNoDescRequestCmd, ShellCompDirectiveError, ShellCompDirectiveNoSpace, ShellCompDirectiveNoFileComp)) 348} 349 350func writePostscript(buf *bytes.Buffer, name string) { 351 name = strings.Replace(name, ":", "__", -1) 352 buf.WriteString(fmt.Sprintf("__start_%s()\n", name)) 353 buf.WriteString(fmt.Sprintf(`{ 354 local cur prev words cword 355 declare -A flaghash 2>/dev/null || : 356 declare -A aliashash 2>/dev/null || : 357 if declare -F _init_completion >/dev/null 2>&1; then 358 _init_completion -s || return 359 else 360 __%[1]s_init_completion -n "=" || return 361 fi 362 363 local c=0 364 local flags=() 365 local two_word_flags=() 366 local local_nonpersistent_flags=() 367 local flags_with_completion=() 368 local flags_completion=() 369 local commands=("%[1]s") 370 local must_have_one_flag=() 371 local must_have_one_noun=() 372 local has_completion_function 373 local last_command 374 local nouns=() 375 376 __%[1]s_handle_word 377} 378 379`, name)) 380 buf.WriteString(fmt.Sprintf(`if [[ $(type -t compopt) = "builtin" ]]; then 381 complete -o default -F __start_%s %s 382else 383 complete -o default -o nospace -F __start_%s %s 384fi 385 386`, name, name, name, name)) 387 buf.WriteString("# ex: ts=4 sw=4 et filetype=sh\n") 388} 389 390func writeCommands(buf *bytes.Buffer, cmd *Command) { 391 buf.WriteString(" commands=()\n") 392 for _, c := range cmd.Commands() { 393 if !c.IsAvailableCommand() || c == cmd.helpCommand { 394 continue 395 } 396 buf.WriteString(fmt.Sprintf(" commands+=(%q)\n", c.Name())) 397 writeCmdAliases(buf, c) 398 } 399 buf.WriteString("\n") 400} 401 402func writeFlagHandler(buf *bytes.Buffer, name string, annotations map[string][]string, cmd *Command) { 403 for key, value := range annotations { 404 switch key { 405 case BashCompFilenameExt: 406 buf.WriteString(fmt.Sprintf(" flags_with_completion+=(%q)\n", name)) 407 408 var ext string 409 if len(value) > 0 { 410 ext = fmt.Sprintf("__%s_handle_filename_extension_flag ", cmd.Root().Name()) + strings.Join(value, "|") 411 } else { 412 ext = "_filedir" 413 } 414 buf.WriteString(fmt.Sprintf(" flags_completion+=(%q)\n", ext)) 415 case BashCompCustom: 416 buf.WriteString(fmt.Sprintf(" flags_with_completion+=(%q)\n", name)) 417 if len(value) > 0 { 418 handlers := strings.Join(value, "; ") 419 buf.WriteString(fmt.Sprintf(" flags_completion+=(%q)\n", handlers)) 420 } else { 421 buf.WriteString(" flags_completion+=(:)\n") 422 } 423 case BashCompSubdirsInDir: 424 buf.WriteString(fmt.Sprintf(" flags_with_completion+=(%q)\n", name)) 425 426 var ext string 427 if len(value) == 1 { 428 ext = fmt.Sprintf("__%s_handle_subdirs_in_dir_flag ", cmd.Root().Name()) + value[0] 429 } else { 430 ext = "_filedir -d" 431 } 432 buf.WriteString(fmt.Sprintf(" flags_completion+=(%q)\n", ext)) 433 } 434 } 435} 436 437func writeShortFlag(buf *bytes.Buffer, flag *pflag.Flag, cmd *Command) { 438 name := flag.Shorthand 439 format := " " 440 if len(flag.NoOptDefVal) == 0 { 441 format += "two_word_" 442 } 443 format += "flags+=(\"-%s\")\n" 444 buf.WriteString(fmt.Sprintf(format, name)) 445 writeFlagHandler(buf, "-"+name, flag.Annotations, cmd) 446} 447 448func writeFlag(buf *bytes.Buffer, flag *pflag.Flag, cmd *Command) { 449 name := flag.Name 450 format := " flags+=(\"--%s" 451 if len(flag.NoOptDefVal) == 0 { 452 format += "=" 453 } 454 format += "\")\n" 455 buf.WriteString(fmt.Sprintf(format, name)) 456 if len(flag.NoOptDefVal) == 0 { 457 format = " two_word_flags+=(\"--%s\")\n" 458 buf.WriteString(fmt.Sprintf(format, name)) 459 } 460 writeFlagHandler(buf, "--"+name, flag.Annotations, cmd) 461} 462 463func writeLocalNonPersistentFlag(buf *bytes.Buffer, flag *pflag.Flag) { 464 name := flag.Name 465 format := " local_nonpersistent_flags+=(\"--%s" 466 if len(flag.NoOptDefVal) == 0 { 467 format += "=" 468 } 469 format += "\")\n" 470 buf.WriteString(fmt.Sprintf(format, name)) 471} 472 473// Setup annotations for go completions for registered flags 474func prepareCustomAnnotationsForFlags(cmd *Command) { 475 for flag := range flagCompletionFunctions { 476 // Make sure the completion script calls the __*_go_custom_completion function for 477 // every registered flag. We need to do this here (and not when the flag was registered 478 // for completion) so that we can know the root command name for the prefix 479 // of __<prefix>_go_custom_completion 480 if flag.Annotations == nil { 481 flag.Annotations = map[string][]string{} 482 } 483 flag.Annotations[BashCompCustom] = []string{fmt.Sprintf("__%[1]s_handle_go_custom_completion", cmd.Root().Name())} 484 } 485} 486 487func writeFlags(buf *bytes.Buffer, cmd *Command) { 488 prepareCustomAnnotationsForFlags(cmd) 489 buf.WriteString(` flags=() 490 two_word_flags=() 491 local_nonpersistent_flags=() 492 flags_with_completion=() 493 flags_completion=() 494 495`) 496 localNonPersistentFlags := cmd.LocalNonPersistentFlags() 497 cmd.NonInheritedFlags().VisitAll(func(flag *pflag.Flag) { 498 if nonCompletableFlag(flag) { 499 return 500 } 501 writeFlag(buf, flag, cmd) 502 if len(flag.Shorthand) > 0 { 503 writeShortFlag(buf, flag, cmd) 504 } 505 if localNonPersistentFlags.Lookup(flag.Name) != nil { 506 writeLocalNonPersistentFlag(buf, flag) 507 } 508 }) 509 cmd.InheritedFlags().VisitAll(func(flag *pflag.Flag) { 510 if nonCompletableFlag(flag) { 511 return 512 } 513 writeFlag(buf, flag, cmd) 514 if len(flag.Shorthand) > 0 { 515 writeShortFlag(buf, flag, cmd) 516 } 517 }) 518 519 buf.WriteString("\n") 520} 521 522func writeRequiredFlag(buf *bytes.Buffer, cmd *Command) { 523 buf.WriteString(" must_have_one_flag=()\n") 524 flags := cmd.NonInheritedFlags() 525 flags.VisitAll(func(flag *pflag.Flag) { 526 if nonCompletableFlag(flag) { 527 return 528 } 529 for key := range flag.Annotations { 530 switch key { 531 case BashCompOneRequiredFlag: 532 format := " must_have_one_flag+=(\"--%s" 533 if flag.Value.Type() != "bool" { 534 format += "=" 535 } 536 format += "\")\n" 537 buf.WriteString(fmt.Sprintf(format, flag.Name)) 538 539 if len(flag.Shorthand) > 0 { 540 buf.WriteString(fmt.Sprintf(" must_have_one_flag+=(\"-%s\")\n", flag.Shorthand)) 541 } 542 } 543 } 544 }) 545} 546 547func writeRequiredNouns(buf *bytes.Buffer, cmd *Command) { 548 buf.WriteString(" must_have_one_noun=()\n") 549 sort.Sort(sort.StringSlice(cmd.ValidArgs)) 550 for _, value := range cmd.ValidArgs { 551 // Remove any description that may be included following a tab character. 552 // Descriptions are not supported by bash completion. 553 value = strings.Split(value, "\t")[0] 554 buf.WriteString(fmt.Sprintf(" must_have_one_noun+=(%q)\n", value)) 555 } 556 if cmd.ValidArgsFunction != nil { 557 buf.WriteString(" has_completion_function=1\n") 558 } 559} 560 561func writeCmdAliases(buf *bytes.Buffer, cmd *Command) { 562 if len(cmd.Aliases) == 0 { 563 return 564 } 565 566 sort.Sort(sort.StringSlice(cmd.Aliases)) 567 568 buf.WriteString(fmt.Sprint(` if [[ -z "${BASH_VERSION}" || "${BASH_VERSINFO[0]}" -gt 3 ]]; then`, "\n")) 569 for _, value := range cmd.Aliases { 570 buf.WriteString(fmt.Sprintf(" command_aliases+=(%q)\n", value)) 571 buf.WriteString(fmt.Sprintf(" aliashash[%q]=%q\n", value, cmd.Name())) 572 } 573 buf.WriteString(` fi`) 574 buf.WriteString("\n") 575} 576func writeArgAliases(buf *bytes.Buffer, cmd *Command) { 577 buf.WriteString(" noun_aliases=()\n") 578 sort.Sort(sort.StringSlice(cmd.ArgAliases)) 579 for _, value := range cmd.ArgAliases { 580 buf.WriteString(fmt.Sprintf(" noun_aliases+=(%q)\n", value)) 581 } 582} 583 584func gen(buf *bytes.Buffer, cmd *Command) { 585 for _, c := range cmd.Commands() { 586 if !c.IsAvailableCommand() || c == cmd.helpCommand { 587 continue 588 } 589 gen(buf, c) 590 } 591 commandName := cmd.CommandPath() 592 commandName = strings.Replace(commandName, " ", "_", -1) 593 commandName = strings.Replace(commandName, ":", "__", -1) 594 595 if cmd.Root() == cmd { 596 buf.WriteString(fmt.Sprintf("_%s_root_command()\n{\n", commandName)) 597 } else { 598 buf.WriteString(fmt.Sprintf("_%s()\n{\n", commandName)) 599 } 600 601 buf.WriteString(fmt.Sprintf(" last_command=%q\n", commandName)) 602 buf.WriteString("\n") 603 buf.WriteString(" command_aliases=()\n") 604 buf.WriteString("\n") 605 606 writeCommands(buf, cmd) 607 writeFlags(buf, cmd) 608 writeRequiredFlag(buf, cmd) 609 writeRequiredNouns(buf, cmd) 610 writeArgAliases(buf, cmd) 611 buf.WriteString("}\n\n") 612} 613 614// GenBashCompletion generates bash completion file and writes to the passed writer. 615func (c *Command) GenBashCompletion(w io.Writer) error { 616 buf := new(bytes.Buffer) 617 writePreamble(buf, c.Name()) 618 if len(c.BashCompletionFunction) > 0 { 619 buf.WriteString(c.BashCompletionFunction + "\n") 620 } 621 gen(buf, c) 622 writePostscript(buf, c.Name()) 623 624 _, err := buf.WriteTo(w) 625 return err 626} 627 628func nonCompletableFlag(flag *pflag.Flag) bool { 629 return flag.Hidden || len(flag.Deprecated) > 0 630} 631 632// GenBashCompletionFile generates bash completion file. 633func (c *Command) GenBashCompletionFile(filename string) error { 634 outFile, err := os.Create(filename) 635 if err != nil { 636 return err 637 } 638 defer outFile.Close() 639 640 return c.GenBashCompletion(outFile) 641} 642