Implement bash completions!

Author: PapaCharlie

bash_completions.go

@@ -15,6 +15,7 @@ import (
 const (
 	BashCompFilenameExt     = "cobra_annotation_bash_completion_filename_extensions"
 	BashCompCustom          = "cobra_annotation_bash_completion_custom"
+	BashCompDynamic         = "cobra_annotation_bash_completion_dynamic"
 	BashCompOneRequiredFlag = "cobra_annotation_bash_completion_one_required_flag"
 	BashCompSubdirsInDir    = "cobra_annotation_bash_completion_subdirs_in_dir"
 )
@@ -272,6 +273,29 @@ __%[1]s_handle_word()
     __%[1]s_handle_word
 }
 
+__%[1]s_handle_dynamic_flag_completion()
+{
+    export COBRA_FLAG_COMPLETION="$1"
+
+    local output
+    if ! output="$(mktemp)" ; then
+        return $?
+    fi
+
+    if ! error_message="$($COMP_LINE > "$output")" ; then
+        local st="$?"
+        echo "$error_message"
+        return "$st"
+    fi
+
+    while read -r -d '' line ; do
+        COMPREPLY+=("$line")
+    done < "$output"
+
+    unset COBRA_FLAG_COMPLETION
+    rm "$output"
+}
+
 `, name))
 }
 
@@ -347,6 +371,9 @@ func writeFlagHandler(buf *bytes.Buffer, name string, annotations map[string][]s
 			} else {
 				buf.WriteString("    flags_completion+=(:)\n")
 			}
+		case BashCompDynamic:
+			buf.WriteString(fmt.Sprintf("    flags_with_completion+=(%q)\n", name))
+			buf.WriteString(fmt.Sprintf("    flags_completion+=(%q)\n", value[0]))
 		case BashCompSubdirsInDir:
 			buf.WriteString(fmt.Sprintf("    flags_with_completion+=(%q)\n", name))
 
@@ -519,6 +546,12 @@ func gen(buf *bytes.Buffer, cmd *Command) {
 
 // GenBashCompletion generates bash completion file and writes to the passed writer.
 func (c *Command) GenBashCompletion(w io.Writer) error {
+	visitAllFlagsWithCompletions(c, func(f *pflag.Flag) {
+		if f.Annotations == nil {
+			f.Annotations = make(map[string][]string)
+		}
+		f.Annotations[BashCompDynamic] = []string{"__" + c.Root().Name() + "_handle_dynamic_flag_completion " + f.Name}
+	})
 	buf := new(bytes.Buffer)
 	writePreamble(buf, c.Name())
 	if len(c.BashCompletionFunction) > 0 {

command.go

@@ -19,6 +19,7 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"log"
 	"os"
 	"path/filepath"
 	"sort"
@@ -883,6 +884,10 @@ func (c *Command) complete(flagName string, a []string) (err error) {
 			c.Flags().AddFlag(f)
 		}
 	})
+	if flagToComplete == nil {
+		log.Panicln(flagName, "is not a known flag")
+	}
+
 	if flagToComplete.Shorthand != "" {
 		c.Flags().StringVarP(&currentCompletionValue, flagName, flagToComplete.Shorthand, "", "")
 	} else {
@@ -939,7 +944,7 @@ func (c *Command) complete(flagName string, a []string) (err error) {
 	}
 
 	for _, v := range values {
-		c.Print(v + "\x00")
+		fmt.Print(v + "\x00")
 	}
 
 	return nil
@@ -1696,3 +1701,16 @@ func (c *Command) updateParentsPflags() {
 		c.parentsPflags.AddFlagSet(parent.PersistentFlags())
 	})
 }
+
+func visitAllFlagsWithCompletions(c *Command, fn func(*flag.Flag)) {
+	filterFunc := func(f *flag.Flag) {
+		if _, ok := c.flagCompletions[f]; ok {
+			fn(f)
+		}
+	}
+	c.Flags().VisitAll(filterFunc)
+	c.PersistentFlags().VisitAll(filterFunc)
+	for _, sc := range c.Commands() {
+		visitAllFlagsWithCompletions(sc, fn)
+	}
+}

shell_completions.go

@@ -93,7 +93,7 @@ type DynamicFlagCompletion func(currentValue string) (suggestedValues []string,
 // RunPreRunsDuringCompletion is set to true. All flags other than the flag currently being completed will be parsed
 // according to their type. The flag being completed is parsed as a raw string with no format requirements
 //
-// Shell Completion compatibility matrix: zsh
+// Shell Completion compatibility matrix: bash, zsh
 func (c *Command) MarkDynamicFlagCompletion(name string, completion DynamicFlagCompletion) error {
 	flag := c.Flag(name)
 	if flag == nil {
@@ -107,9 +107,6 @@ func (c *Command) MarkDynamicFlagCompletion(name string, completion DynamicFlagC
 		c.flagCompletions = make(map[*pflag.Flag]DynamicFlagCompletion)
 	}
 	c.flagCompletions[flag] = completion
-	if flag.Annotations == nil {
-		flag.Annotations = map[string][]string{}
-	}
-	flag.Annotations[zshCompDynamicCompletion] = []string{zshCompGenFlagCompletionFuncName(c)}
+
 	return nil
 }

zsh_completions.go

@@ -26,9 +26,8 @@ var (
 		"extractFlags":                        zshCompExtractFlag,
 		"genFlagEntryForZshArguments":         zshCompGenFlagEntryForArguments,
 		"extractArgsCompletions":              zshCompExtractArgumentCompletionHintsForRendering,
-		"genZshFlagDynamicCompletionFuncName": zshCompGenFlagCompletionFuncName,
+		"genZshFlagDynamicCompletionFuncName": zshCompGenDynamicFlagCompletionFuncName,
 		"hasDynamicCompletions":               zshCompHasDynamicCompletions,
-		"flagCompletionsEnvVar":               func() string { return FlagCompletionEnvVar },
 	}
 	zshCompletionText = `
 {{/* should accept Command (that contains subcommands) as parameter */}}
@@ -86,16 +85,27 @@ function {{genZshFuncName .}} {
 {{if hasDynamicCompletions . -}}
 function {{genZshFlagDynamicCompletionFuncName .}} {
   export COBRA_FLAG_COMPLETION="$1"
-  if suggestions="$("$words[@]" 2>&1)" ; then
-    local -a args
-    while read -d $'\0' line ; do
-      args+="$line"
-    done <<< "$suggestions"
-    _values "$1" "$args[@]"
-  else
-    _message "Exception occurred during completion: $suggestions"
+
+  local output
+  if ! output="$(mktemp)" ; then
+    return $?
+  fi
+
+  if ! error_message="$("${tokens[@]}" 2>&1 > "$output")" ; then
+    local st="$?"
+    _message "Exception occurred during completion: $error_message"
+    return "$st"
   fi
+
+  local -a args
+  while read -r -d '' line ; do
+    args+="$line"
+  done < "$output"
+
+  _values "$1" "$args[@]"
+
   unset COBRA_FLAG_COMPLETION
+  rm "$output"
 }{{- end}}
 
 {{template "selectCmdTemplate" .}}
@@ -131,6 +141,12 @@ func (c *Command) GenZshCompletionFile(filename string) error {
 // writer. The completion always run on the root command regardless of the
 // command it was called from.
 func (c *Command) GenZshCompletion(w io.Writer) error {
+	visitAllFlagsWithCompletions(c, func(f *pflag.Flag) {
+		if f.Annotations == nil {
+			f.Annotations = make(map[string][]string)
+		}
+		f.Annotations[zshCompDynamicCompletion] = []string{zshCompGenDynamicFlagCompletionFuncName(c)}
+	})
 	tmpl, err := template.New("Main").Funcs(zshCompFuncMap).Parse(zshCompletionText)
 	if err != nil {
 		return fmt.Errorf("error creating zsh completion template: %v", err)
@@ -269,7 +285,7 @@ func zshCompGenFuncName(c *Command) string {
 	return "_" + c.Name()
 }
 
-func zshCompGenFlagCompletionFuncName(c *Command) string {
+func zshCompGenDynamicFlagCompletionFuncName(c *Command) string {
 	return "_" + c.Root().Name() + "-flag-completion"
 }