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