feat(comp): Dynamic completion of flags in Go

Conflicts:
	cmd/helm/completion/complete.go
	cmd/helm/root.go

Conflicts:
	cmd/helm/root.go

Signed-off-by: Marc Khouzam <marc.khouzam@montreal.ca>

cmd/helm/flags.go

@@ -23,6 +23,7 @@ import (
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 
+	"helm.sh/helm/v3/internal/completion"
 	"helm.sh/helm/v3/pkg/action"
 	"helm.sh/helm/v3/pkg/cli/output"
 	"helm.sh/helm/v3/pkg/cli/values"
@@ -52,10 +53,20 @@ func addChartPathOptionsFlags(f *pflag.FlagSet, c *action.ChartPathOptions) {
 // bindOutputFlag will add the output flag to the given command and bind the
 // value to the given format pointer
 func bindOutputFlag(cmd *cobra.Command, varRef *output.Format) {
-	cmd.Flags().VarP(newOutputValue(output.Table, varRef), outputFlag, "o",
+	f := cmd.Flags()
+	f.VarP(newOutputValue(output.Table, varRef), outputFlag, "o",
 		fmt.Sprintf("prints the output in the specified format. Allowed values: %s", strings.Join(output.Formats(), ", ")))
-	// Setup shell completion for the flag
+
-	cmd.MarkFlagCustom(outputFlag, "__helm_output_options")
+	flag := f.Lookup(outputFlag)
+	completion.RegisterFlagCompletionFunc(flag, func(cmd *cobra.Command, args []string, toComplete string) ([]string, completion.BashCompDirective) {
+		var formatNames []string
+		for _, format := range output.Formats() {
+			if strings.HasPrefix(format, toComplete) {
+				formatNames = append(formatNames, format)
+			}
+		}
+		return formatNames, completion.BashCompDirectiveDefault
+	})
 }
 
 type outputValue output.Format

cmd/helm/root.go

@@ -23,51 +23,16 @@ import (
 
 	"github.com/spf13/cobra"
 
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
 	"helm.sh/helm/v3/cmd/helm/require"
 	"helm.sh/helm/v3/internal/completion"
 	"helm.sh/helm/v3/internal/experimental/registry"
 	"helm.sh/helm/v3/pkg/action"
-	"helm.sh/helm/v3/pkg/cli/output"
 )
 
 const (
 	bashCompletionFunc = `
-__helm_override_flag_list=(--kubeconfig --kube-context --namespace -n)
-__helm_override_flags()
-{
-    local ${__helm_override_flag_list[*]##*-} two_word_of of var
-    for w in "${words[@]}"; do
-        if [ -n "${two_word_of}" ]; then
-            eval "${two_word_of##*-}=\"${two_word_of}=\${w}\""
-            two_word_of=
-            continue
-        fi
-        for of in "${__helm_override_flag_list[@]}"; do
-            case "${w}" in
-                ${of}=*)
-                    eval "${of##*-}=\"${w}\""
-                    ;;
-                ${of})
-                    two_word_of="${of}"
-                    ;;
-            esac
-        done
-    done
-    for var in "${__helm_override_flag_list[@]##*-}"; do
-        if eval "test -n \"\$${var}\""; then
-            eval "echo \${${var}}"
-        fi
-    done
-}
-
-__helm_override_flags_to_kubectl_flags()
-{
-    # --kubeconfig, -n, --namespace stay the same for kubectl
-    # --kube-context becomes --context for kubectl
-    __helm_debug "${FUNCNAME[0]}: flags to convert: $1"
-    echo "$1" | \sed s/kube-context/context/
-}
-
 __helm_get_contexts()
 {
     __helm_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}"
@@ -78,36 +43,19 @@ __helm_get_contexts()
     fi
 }
 
-__helm_get_namespaces()
-{
-    __helm_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}"
-    local template out
-    template="{{ range .items  }}{{ .metadata.name }} {{ end }}"
-
-    flags=$(__helm_override_flags_to_kubectl_flags "$(__helm_override_flags)")
-    __helm_debug "${FUNCNAME[0]}: override flags for kubectl are: $flags"
-
-    # Must use eval in case the flags contain a variable such as $HOME
-    if out=$(eval kubectl get ${flags} -o template --template=\"${template}\" namespace 2>/dev/null); then
-        COMPREPLY+=( $( compgen -W "${out[*]}" -- "$cur" ) )
-    fi
-}
-
-__helm_output_options()
-{
-    __helm_debug "${FUNCNAME[0]}: c is $c words[c] is ${words[c]}"
-    COMPREPLY+=( $( compgen -W "%[1]s" -- "$cur" ) )
-}
-
 __helm_custom_func()
 {
     __helm_debug "${FUNCNAME[0]}: c is $c, words[@] is ${words[@]}, #words[@] is ${#words[@]}"
     __helm_debug "${FUNCNAME[0]}: cur is ${cur}, cword is ${cword}, words is ${words}"
 
-    local out requestComp
+    local out requestComp lastParam lastChar
-    requestComp="${words[0]} %[2]s ${words[@]:1}"
+    requestComp="${words[0]} %[1]s ${words[@]:1}"
+
+    lastParam=${words[$((${#words[@]}-1))]}
+    lastChar=${lastParam:$((${#lastParam}-1)):1}
+    __helm_debug "${FUNCNAME[0]}: lastParam ${lastParam}, lastChar ${lastChar}"
 
-    if [ -z "${cur}" ]; then
+    if [ -z "${cur}" ] && [ "${lastChar}" != "=" ]; then
         # If the last parameter is complete (there is a space following it)
         # We add an extra empty parameter so we can indicate this to the go method.
         __helm_debug "${FUNCNAME[0]}: Adding extra empty parameter"
@@ -119,15 +67,15 @@ __helm_custom_func()
     out=$(eval ${requestComp} 2>/dev/null)
     directive=$?
 
-    if [ $((${directive} & %[3]d)) -ne 0 ]; then
+    if [ $((${directive} & %[2]d)) -ne 0 ]; then
         __helm_debug "${FUNCNAME[0]}: received error, completion failed"
     else
-        if [ $((${directive} & %[4]d)) -ne 0 ]; then
+        if [ $((${directive} & %[3]d)) -ne 0 ]; then
             if [[ $(type -t compopt) = "builtin" ]]; then
                 compopt -o nospace
             fi
         fi
-        if [ $((${directive} & %[5]d)) -ne 0 ]; then
+        if [ $((${directive} & %[4]d)) -ne 0 ]; then
             if [[ $(type -t compopt) = "builtin" ]]; then
                 compopt +o default
             fi
@@ -145,7 +93,8 @@ var (
 	// Mapping of global flags that can have dynamic completion and the
 	// completion function to be used.
 	bashCompletionFlags = map[string]string{
-		"namespace":    "__helm_get_namespaces",
+		// Cannot convert the kube-context flag to Go completion yet because
+		// an incomplete kube-context will make actionConfig.Init() fail at the very start
 		"kube-context": "__helm_get_contexts",
 	}
 )
@@ -198,7 +147,6 @@ func newRootCmd(actionConfig *action.Configuration, out io.Writer, args []string
 		Args:         require.NoArgs,
 		BashCompletionFunction: fmt.Sprintf(
 			bashCompletionFunc,
-			strings.Join(output.Formats(), " "),
 			completion.CompRequestCmd,
 			completion.BashCompDirectiveError,
 			completion.BashCompDirectiveNoSpace,
@@ -208,6 +156,29 @@ func newRootCmd(actionConfig *action.Configuration, out io.Writer, args []string
 
 	settings.AddFlags(flags)
 
+	flag := flags.Lookup("namespace")
+	// Setup shell completion for the namespace flag
+	completion.RegisterFlagCompletionFunc(flag, func(cmd *cobra.Command, args []string, toComplete string) ([]string, completion.BashCompDirective) {
+		if client, err := actionConfig.KubernetesClientSet(); err == nil {
+			// Choose a long enough timeout that the user notices somethings is not working
+			// but short enough that the user is not made to wait very long
+			to := int64(3)
+			completion.CompDebugln(fmt.Sprintf("About to call kube client for namespaces with timeout of: %d", to))
+
+			nsNames := []string{}
+			// TODO can we request only the namespace with request.prefix?
+			if namespaces, err := client.CoreV1().Namespaces().List(metav1.ListOptions{TimeoutSeconds: &to}); err == nil {
+				for _, ns := range namespaces.Items {
+					if strings.HasPrefix(ns.Name, toComplete) {
+						nsNames = append(nsNames, ns.Name)
+					}
+				}
+				return nsNames, completion.BashCompDirectiveNoFileComp
+			}
+		}
+		return nil, completion.BashCompDirectiveDefault
+	})
+
 	// We can safely ignore any errors that flags.Parse encounters since
 	// those errors will be caught later during the call to cmd.Execution.
 	// This call is required to gather configuration information prior to

internal/completion/complete.go

@@ -22,6 +22,7 @@ import (
 	"strings"
 
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
 
 	"helm.sh/helm/v3/cmd/helm/require"
 	"helm.sh/helm/v3/pkg/cli"
@@ -36,8 +37,8 @@ import (
 // Used by the shell completion script.
 const CompRequestCmd = "__complete"
 
-// Global map allowing to find completion functions for commands.
+// Global map allowing to find completion functions for commands or flags.
-var validArgsFunctions = map[*cobra.Command]func(cmd *cobra.Command, args []string, toComplete string) ([]string, BashCompDirective){}
+var validArgsFunctions = map[interface{}]func(cmd *cobra.Command, args []string, toComplete string) ([]string, BashCompDirective){}
 
 // BashCompDirective is a bit map representing the different behaviors the shell
 // can be instructed to have once completions have been provided.
@@ -69,6 +70,21 @@ func RegisterValidArgsFunc(cmd *cobra.Command, f func(cmd *cobra.Command, args [
 	validArgsFunctions[cmd] = f
 }
 
+// RegisterFlagCompletionFunc should be called to register a function to provide completion for a flag
+func RegisterFlagCompletionFunc(flag *pflag.Flag, f func(cmd *cobra.Command, args []string, toComplete string) ([]string, BashCompDirective)) {
+	if _, exists := validArgsFunctions[flag]; exists {
+		log.Fatal(fmt.Sprintf("RegisterFlagCompletionFunc: flag '%s' already registered", flag.Name))
+	}
+	validArgsFunctions[flag] = f
+
+	// Make sure the completion script call the __helm_custom_func for the registered flag.
+	// This is essential to make the = form work. E.g., helm -n=<TAB> or helm status --output=<TAB>
+	if flag.Annotations == nil {
+		flag.Annotations = map[string][]string{}
+	}
+	flag.Annotations[cobra.BashCompCustom] = []string{"__helm_custom_func"}
+}
+
 var debug = true
 
 // Returns a string listing the different directive enabled in the specified parameter
@@ -101,15 +117,14 @@ func NewCompleteCmd(settings *cli.EnvSettings) *cobra.Command {
 		DisableFlagsInUseLine: true,
 		Hidden:                true,
 		DisableFlagParsing:    true,
-		Args:                  require.MinimumNArgs(2),
+		Args:                  require.MinimumNArgs(1),
 		Short:                 "Request shell completion choices for the specified command-line",
 		Long: fmt.Sprintf("%s is a special command that is used by the shell completion logic\n%s",
 			CompRequestCmd, "to request completion choices for the specified command-line."),
 		Run: func(cmd *cobra.Command, args []string) {
 			CompDebugln(fmt.Sprintf("%s was called with args %v", cmd.Name(), args))
 
-			trimmedArgs := args[:len(args)-1]
+			flag, trimmedArgs, toComplete := checkIfFlagCompletion(cmd.Root(), args[:len(args)-1], args[len(args)-1])
-			toComplete := args[len(args)-1]
 
 			// Find the real command for which completion must be performed
 			finalCmd, finalArgs, err := cmd.Root().Find(trimmedArgs)
@@ -128,10 +143,20 @@ func NewCompleteCmd(settings *cli.EnvSettings) *cobra.Command {
 			argsWoFlags := finalCmd.Flags().Args()
 			CompDebugln(fmt.Sprintf("Args without flags are '%v' with length %d", argsWoFlags, len(argsWoFlags)))
 
-			// Find completion function for the command
+			var key interface{}
-			completionFn, ok := validArgsFunctions[finalCmd]
+			var keyStr string
+			if flag != nil {
+				key = flag
+				keyStr = flag.Name
+			} else {
+				key = finalCmd
+				keyStr = finalCmd.Name()
+			}
+
+			// Find completion function for the flag or command
+			completionFn, ok := validArgsFunctions[key]
 			if !ok {
-				CompErrorln(fmt.Sprintf("Dynamic completion not supported/needed for flag or command: %s", finalCmd.Name()))
+				CompErrorln(fmt.Sprintf("Dynamic completion not supported/needed for flag or command: %s", keyStr))
 				return
 			}
 
@@ -150,6 +175,96 @@ func NewCompleteCmd(settings *cli.EnvSettings) *cobra.Command {
 	}
 }
 
+func isFlag(arg string) bool {
+	return len(arg) > 0 && arg[0] == '-'
+}
+
+func checkIfFlagCompletion(rootCmd *cobra.Command, args []string, lastArg string) (*pflag.Flag, []string, string) {
+	var flagName string
+	trimmedArgs := args
+	flagWithEqual := false
+	if isFlag(lastArg) {
+		if index := strings.Index(lastArg, "="); index >= 0 {
+			flagName = strings.TrimLeft(lastArg[:index], "-")
+			lastArg = lastArg[index+1:]
+			flagWithEqual = true
+		} else {
+			CompErrorln("Unexpected completion request for flag")
+			os.Exit(int(BashCompDirectiveError))
+		}
+	}
+
+	if len(flagName) == 0 {
+		if len(args) > 0 {
+			prevArg := args[len(args)-1]
+			if isFlag(prevArg) {
+				// If the flag contains an = it means it has already been fully processed
+				if index := strings.Index(prevArg, "="); index < 0 {
+					flagName = strings.TrimLeft(prevArg, "-")
+
+					// Remove the uncompleted flag or else Cobra could complain about
+					// an invalid value for that flag e.g., helm status --output j<TAB>
+					trimmedArgs = args[:len(args)-1]
+				}
+			}
+		}
+	}
+
+	if len(flagName) == 0 {
+		// Not doing flag completion
+		return nil, trimmedArgs, lastArg
+	}
+
+	// Find the real command for which completion must be performed
+	finalCmd, _, err := rootCmd.Find(trimmedArgs)
+	if err != nil {
+		// Unable to find the real command. E.g., helm invalidCmd <TAB>
+		os.Exit(int(BashCompDirectiveError))
+	}
+
+	CompDebugln(fmt.Sprintf("checkIfFlagCompletion: found final command '%s'", finalCmd.Name()))
+
+	flag := findFlag(finalCmd, flagName)
+	if flag == nil {
+		// Flag not supported by this command, nothing to complete
+		CompDebugln(fmt.Sprintf("Subcommand '%s' does not support flag '%s'", finalCmd.Name(), flagName))
+		os.Exit(int(BashCompDirectiveNoFileComp))
+	}
+
+	if !flagWithEqual {
+		if len(flag.NoOptDefVal) != 0 {
+			// We had assumed dealing with a two-word flag but the flag is a boolean flag.
+			// In that case, there is no value following it, so we are not really doing flag completion.
+			// Reset everything to do argument completion.
+			trimmedArgs = args
+			flag = nil
+		}
+	}
+
+	return flag, trimmedArgs, lastArg
+}
+
+func findFlag(cmd *cobra.Command, name string) *pflag.Flag {
+	flagSet := cmd.Flags()
+	if len(name) == 1 {
+		// First convert the short flag into a long flag
+		// as the cmd.Flag() search only accepts long flags
+		if short := flagSet.ShorthandLookup(name); short != nil {
+			CompDebugln(fmt.Sprintf("checkIfFlagCompletion: found flag '%s' which we will change to '%s'", name, short.Name))
+			name = short.Name
+		} else {
+			set := cmd.InheritedFlags()
+			if short = set.ShorthandLookup(name); short != nil {
+				CompDebugln(fmt.Sprintf("checkIfFlagCompletion: found inherited flag '%s' which we will change to '%s'", name, short.Name))
+				name = short.Name
+			} else {
+				return nil
+			}
+		}
+	}
+	return cmd.Flag(name)
+}
+
 // CompDebug prints the specified string to the same file as where the
 // completion script prints its logs.
 // Note that completion printouts should never be on stdout as they would