Fix zsh completion handling of nospace and file completion

[marckhouzam]

Fixes #1211
Fixes #1212

This PR improves the zsh completion script to allow it to properly handle ShellDirectiveCompletionNoSpace and file completion all the time. Specifically, there was an issue where when ValidArgsFunction would always return completions even if they don't match what the user types, the script would not behave right. Please refer to #1211 and #1212 for specific details.

So this PR no longer counts the number of completions it receives, but instead relies on the zsh _describe function and its return code to handle things cleanly.

[marckhouzam]

@jharshman If you can find time to review this before doing the next release it would avoid releasing this bug. Thank you in advance.

[marckhouzam]

This change is unrelated but I just wanted to get it in at the same time.
It allows a program to use extra directives if it wants to provide its own shell script, instead of using Cobra's.

[marckhouzam]

Here is how to test this specific fix. An extensive test program is included at the bottom.

Test 1 - ShellCompDirectiveNoSpace not respected
Before the PR:

$ zsh
$ source <(./testprog completion zsh)
$ compdef _testprog testprog

$ ./testprog nospace compo<TAB>
# and
$ ./testprog --nospace compo<TAB>

will complete to:

$ ./testprog nospace compost<SPACE>
# and
$ ./testprog --nospace compost<SPACE>

the <SPACE> is wrong at the end of the completion because ValidArgsFunction requested ShellCompDirectiveNoSpace, as can be seen by:

$ ./testprog __complete nospace compo
complete	description1
compost	description2
:2
Completion ended with directive: ShellCompDirectiveNoSpace

Testing with this PR in the exact same way will not output the <SPACE> after the word compost.

Test 2 - ShellCompDirectiveNoSpace not properly handling descriptions

Before the PR:

$ zsh
$ source <(./testprog completion zsh)
$ compdef _testprog testprog

$ ./testprog nospace first compo<TAB>

will complete to:

$ ./testprog nospace first compost:description2

where we can see that the description is mistakenly printed.

Testing with this PR in the exact same way will not output the description after the word compost.

Test 3 - File completion not always done when requested

Before the PR:

$ zsh
$ source <(./testprog completion zsh)
$ compdef _testprog testprog
$ touch myfile

$ ./testprog withfile myfi<TAB>

will not perform file completion and will not show myfile although it should have.

Testing with this PR in the exact same way will trigger file completion and will show myfile as expected.

Extensive test program to execute the above tests, but also any regression tests:

package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/spf13/cobra"
)

var completionNoDesc bool

var completionCmd = &cobra.Command{
	Use:   "completion [bashzsh|fish|powershell]",
	Short: "Generate completion script",
	Long: `To load completions:

Bash:

$ source <(testprog completion bash)

# To load completions for each session, execute once:
Linux:
  $ testprog completion bash > /etc/bash_completion.d/testprog
MacOS:
  $ testprog completion bash > /usr/local/etc/bash_completion.d/testprog

Zsh:

$ source <(testprog completion zsh)
$ compdef _testprog testprog

# To load completions for each session, execute once:
$ testprog completion zsh > "${fpath[1]}/_testprog"

Fish:

$ testprog completion fish | source

# To load completions for each session, execute once:
$ testprog completion fish > ~/.config/fish/completions/testprog.fish

PowerShell:

PS C:\> testprog completion powershell | Out-String | Invoke-Expression

To load completions for every new session, add the output of the above command
to your powershell profile.
`,
	DisableFlagsInUseLine: true,
	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		return []string{"bash\tBash completion", "zsh", "fish", "powershell"}, cobra.ShellCompDirectiveNoFileComp
	},
	Args: cobra.ExactValidArgs(1),
	Run: func(cmd *cobra.Command, args []string) {
		switch args[0] {
		case "bash":
			cmd.Root().GenBashCompletion(os.Stdout)
		case "zsh":
			if !completionNoDesc {
				cmd.Root().GenZshCompletion(os.Stdout)
			} else {
				cmd.Root().GenZshCompletionNoDesc(os.Stdout)
			}
		case "fish":
			cmd.Root().GenFishCompletion(os.Stdout, !completionNoDesc)
		case "powershell":
			cmd.Root().GenPowerShellCompletion(os.Stdout)
		}
	},
}

var rootCmd = &cobra.Command{
	Use: "testprog",
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("rootCmd called")
	},
}

var vargsCmd = &cobra.Command{
	Use:        "vargs",
	Short:      "Cmd with ValidArgs AND a subcmd",
	ValidArgs:  []string{"one", "two", "three"},
	ArgAliases: []string{"un", "deux", "trois"},
	Args:       cobra.MinimumNArgs(1),
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("vargsCmd called")
	},
}

var childVargsCmd = &cobra.Command{
	Use:   "childVargs",
	Short: "Child command",
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("childVargs called")
	},
}

var aliasCmd = &cobra.Command{
	Use:     "cmdWithAlias",
	Short:   "Command with aliases",
	Aliases: []string{"aliasCmd", "cmdAlias"},
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("cmdWithAlias called")
	},
}

// A command with no short description
var nodescriptionCmd = &cobra.Command{
	Use: "nodesc",
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("nodesc called")
	},
}

// File filtering for the first argument
// Replaces MarkZshCompPositionalArgumentFile(1, string{"*.log", "*.txt"})
var filterFileCmd = &cobra.Command{
	Use:   "filterFile",
	Short: "Only list file of type [*.log,*.txt]",
	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		if len(args) == 0 {
			return []string{"log", "txt"}, cobra.ShellCompDirectiveFilterFileExt
		}
		// No more expected arguments, so turn off file completion
		return nil, cobra.ShellCompDirectiveNoFileComp
	},
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("filterFile called")
	},
}

// Turn off file completion for all arguments
var noFileCmd = &cobra.Command{
	Use:   "nofile",
	Short: "No file completion done",
	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		return nil, cobra.ShellCompDirectiveNoFileComp
	},
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("nofile called")
	},
}

// Full file completion for arguments
var fullFileCmd = &cobra.Command{
	Use:   "fullFile [filename]",
	Short: "Full file completion",
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("fullFile called")
	},
}

// Only allow directory name completion and only for the first argument
var dirOnlyCmd = &cobra.Command{
	Use:   "dirOnly [dirname]",
	Short: "Only list directories in completion",
	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		if len(args) == 0 {
			return nil, cobra.ShellCompDirectiveFilterDirs
		}
		// No more expected arguments, so turn off file completion
		return nil, cobra.ShellCompDirectiveNoFileComp
	},
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("dirOnly called")
	},
}

// Only allow directory name from within a specified directory and only for the first argument
var subdirCmd = &cobra.Command{
	Use:   "subdir [sub-dirname]",
	Short: "Only list directories within themes/",
	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		if len(args) == 0 {
			return []string{"themes"}, cobra.ShellCompDirectiveFilterDirs | cobra.ShellCompDirectiveNoFileComp
		}
		// No more expected arguments, so turn off file completion
		return nil, cobra.ShellCompDirectiveNoFileComp
	},
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("subdir called")
	},
}

// Test file completion when returning other completions
var compAndFileCmd = &cobra.Command{
	Use:   "compAndFile [sun|moon|<filename>]",
	Short: "Provide some completions and file completion when prefix does not match",
	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		if len(args) != 0 {
			return []string{"nginx", "thanos"}, cobra.ShellCompDirectiveNoFileComp
		}
		var completions []string
		for _, comp := range []string{"install", "uninstall"} {
			completions = append(completions, comp)
		}
		return completions, cobra.ShellCompDirectiveDefault
	},
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("compAndFile called")
	},
}

var colonCmd = &cobra.Command{
	Use:       "cmdColonArgs",
	Short:     "Command with colon args",
	ValidArgs: []string{"first:ONE", "first:SECOND"},
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("colonCmd called")
	},
}

var longDescCmd = &cobra.Command{
	Use:       "cmdLongDesc12345678901234567890123456789012345678901234567890",
	Short:     `We have here a quite long description which should go over the length of the shell width.  What should we do with such a long description? zsh and fish handle it nicely, so we should prepare bash to handle it in the same fashion`,
	ValidArgs: []string{"first:ONE", "first:SECOND"},
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("colonCmd called")
	},
}

var specialCharDescCmd = &cobra.Command{
	Use:   "specialChar",
	Short: "Description contains chars like ` and maybe others",
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("specialChar cmd called")
	},
}

var noSpaceCmd = &cobra.Command{
	Use:   "nospace",
	Short: "followed by completion with noSpace, even for file completion",
	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		comps := []string{"complete\tdescription1", "compost\tdescription2"}
		if len(args) == 0 {
			// Allow file completion with ShellCompDirectiveNoSpace
			return comps, cobra.ShellCompDirectiveNoSpace
		}

		var finalComps []string
		for _, comp := range comps {
			if strings.HasPrefix(comp, toComplete) {
				finalComps = append(finalComps, comp)
			}
		}
		return finalComps, cobra.ShellCompDirectiveNoSpace
	},
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("nospace cmd called")
	},
}

var sentenceCmd = &cobra.Command{
	Use:   "sentence",
	Short: "completion has spaces",
	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		// Allow file completion with ShellCompDirectiveNoSpace
		return []string{"comp is a sentence\tFirst comp", "and another sentence\tsecondComp"}, cobra.ShellCompDirectiveNoFileComp
	},
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("sentence cmd called")
	},
}

var withfileCmd = &cobra.Command{
	Use:   "withfile",
	Short: "completion always returned and has file completion",
	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		// We do not consider the toComplete prefix and always return a completion
		return []string{"why"}, cobra.ShellCompDirectiveDefault
	},
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("withfile cmd called")
	},
}

func setFlags() {
	rootCmd.Flags().String("theme", "", "theme to use (located in /themes/THEMENAME/)")
	rootCmd.Flags().SetAnnotation("theme", cobra.BashCompSubdirsInDir, []string{"themes"})

	rootCmd.Flags().String("theme2", "", "theme to use (located in /themes/THEMENAME/)")
	rootCmd.RegisterFlagCompletionFunc("theme2", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		return []string{"themes"}, cobra.ShellCompDirectiveFilterDirs
	})

	rootCmd.Flags().String("customComp", "", "test custom comp for flags")
	rootCmd.RegisterFlagCompletionFunc("customComp", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		return []string{"firstComp\tthe first value", "secondComp\tthe second value", "forthComp\tthe forth value"}, cobra.ShellCompDirectiveNoFileComp
	})

	rootCmd.Flags().StringP("file", "f", "", "list files")
	rootCmd.Flags().String("longdesc", "", "before newline\nafter newline")

	rootCmd.Flags().BoolP("bool", "b", false, "bool flag")
	rootCmd.Flags().BoolSliceP("bslice", "B", nil, "bool slice flag")

	rootCmd.PersistentFlags().StringP("persistent", "p", "", "persistent flag")

	// rootCmd.Flags().String("required", "", "required flag")
	// rootCmd.MarkFlagRequired("required")

	vargsCmd.Flags().StringP("nonpersistent", "n", "", "non persistent local flag")

	// If there is only a recommended value before we reach the equal sign or the colon, then the completion works as expected.
	// If the values differ afterwards, completion fails in bash, but works in fish and zsh
	rootCmd.Flags().String("equalSignWorks", "", "test custom comp with equal sign")
	rootCmd.RegisterFlagCompletionFunc("equalSignWorks", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		return []string{"first=Comp\tthe first value", "second=Comp\tthe second value", "forth=Comp\tthe forth value"}, cobra.ShellCompDirectiveNoFileComp
	})

	rootCmd.Flags().String("equalSignWorksNot", "", "test custom comp with equal sign")
	rootCmd.RegisterFlagCompletionFunc("equalSignWorksNot", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		return []string{"first=Comp1\tthe first value", "first=Comp2\tthe second value", "first=Comp3\tthe forth value"}, cobra.ShellCompDirectiveNoFileComp
	})

	rootCmd.Flags().String("colonWorks", "", "test custom comp with colon one value works")
	rootCmd.RegisterFlagCompletionFunc("colonWorks", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		return []string{"first:ONE\tthe first value", "second\tthe second value", "forth\tthe forth value"}, cobra.ShellCompDirectiveNoFileComp
	})

	rootCmd.Flags().String("colonWorksNot", "", "test custom comp with colon doesnt work")
	rootCmd.RegisterFlagCompletionFunc("colonWorksNot", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		return []string{"first:ONE\tthe first value", "first:SECOND\tthe second value", "second:Comp1\tthe forth value"}, cobra.ShellCompDirectiveNoFileComp
	})

	rootCmd.Flags().String("nospace", "", "test custom comp with nospace directive")
	rootCmd.RegisterFlagCompletionFunc("nospace", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
		return []string{"complete\tFirst comp", "compost\tsecondComp"}, cobra.ShellCompDirectiveNoSpace
	})

	rootCmd.Flags().StringArray("array", []string{"allo", "toi"}, "string array flag")
}

func main() {
	completionCmd.Flags().BoolVar(
		&completionNoDesc,
		"no-descriptions", false,
		"disable completion description for shells that support it")
	rootCmd.AddCommand(completionCmd)

	rootCmd.AddCommand(
		filterFileCmd,
		noFileCmd,
		fullFileCmd,
		dirOnlyCmd,
		subdirCmd,
		compAndFileCmd,
		nodescriptionCmd,
		vargsCmd,
		colonCmd,
		aliasCmd,
		longDescCmd,
		specialCharDescCmd,
		noSpaceCmd,
		sentenceCmd,
		withfileCmd,
	)

	vargsCmd.AddCommand(childVargsCmd)
	setFlags()
	rootCmd.Execute()
}

[marckhouzam]

@jharshman would you be able to find time to review this one?

Helm has moved to Cobra 1.1.1 and has started using native zsh completion. So, for the first time, Helm provides completion descriptions for zsh 🎉 . I'm currently working to also add descriptions to Helm's custom completions (beyond the ones for commands and flags provided by Cobra).

However, in one case ShellCompDirectiveNoSpace is used, and the bug this PR fixes is triggered and the description is added as part of the actual description.

If you can find some time, it would be appreciated. Thanks in advance.

P.S. FYI besides helm, kubectl has also moved to Cobra 1.1.1 and work is on-going to move all its completions to Go completions, and then move to native zsh completions.

[marckhouzam]

The force push was just a rebase.
Anyone up to reviewing this one? I'm hoping to have it in the next Cobra release so helm can use it.
I realize it is not easy to understand the completion script logic, so please don't hesitate to ask for clarifications.

[github-actions]

This PR is being marked as stale due to a long period of inactivity

[marckhouzam]

Waiting for review

[jpmcb]

LGTM! 🚀