From 0745e5576205cd05fada5bb2e01500d97e1a04be Mon Sep 17 00:00:00 2001
From: Sebastiaan van Stijn <thaJeztah@users.noreply.github.com>
Date: Sat, 28 Dec 2024 23:41:43 +0100
Subject: [PATCH] completion: improve detection for flags that accept multiple
 values (#2210)

The completion code attempts to detect whether a flag can be specified
more than once, and therefore should provide completion even if already
set.

Currently, this code depends on conventions used in the pflag package,
which uses an "Array" or "Slice" suffix or for some types a "stringTo"
prefix.

Cobra allows custom value types to be used, which may not use the same
convention for naming, and therefore currently aren't detected to allow
multiple values.

The pflag module defines a [SliceValue] interface, which is implemented
by the Slice and Array value types it provides (unfortunately, it's not
currently implemented by the "stringTo" values).

This patch adds a reduced interface based on the [SliceValue] interface
mentioned above to allow detecting Value-types that accept multiple values.
Custom types can implement this interface to make completion work for
those values.

I deliberately used a reduced interface to keep the requirements for this
detection as low as possible, without enforcing the other methods defined
in the interface (Append, Replace) which may not apply to all custom types.

Future improvements can likely still be made, considering either implementing
the SliceValue interface for the "stringTo" values or defining a separate
"MapValue" interface for those types.

Possibly providing the reduced interface as part of the pflag module and
to export it.

[SliceValue]: https://github.com/spf13/pflag/blob/d5e0c0615acee7028e1e2740a11102313be88de1/flag.go#L193-L203

Signed-off-by: Sebastiaan van Stijn <github@gone.nl>
Signed-off-by: Marc Khouzam <marc.khouzam@gmail.com>
Co-authored-by: Marc Khouzam <marc.khouzam@gmail.com>
---
 completions.go      | 15 +++++++++++++--
 completions_test.go | 35 ++++++++++++++++++++++++++++++++++-
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/completions.go b/completions.go
index 8fccdaf2..0862d3f6 100644
--- a/completions.go
+++ b/completions.go
@@ -270,6 +270,14 @@ func (c *Command) initCompleteCmd(args []string) {
 	}
 }
 
+// SliceValue is a reduced version of [pflag.SliceValue]. It is used to detect
+// flags that accept multiple values and therefore can provide completion
+// multiple times.
+type SliceValue interface {
+	// GetSlice returns the flag value list as an array of strings.
+	GetSlice() []string
+}
+
 func (c *Command) getCompletions(args []string) (*Command, []string, ShellCompDirective, error) {
 	// The last argument, which is not completely typed by the user,
 	// should not be part of the list of arguments
@@ -399,10 +407,13 @@ func (c *Command) getCompletions(args []string) (*Command, []string, ShellCompDi
 		// If we have not found any required flags, only then can we show regular flags
 		if len(completions) == 0 {
 			doCompleteFlags := func(flag *pflag.Flag) {
-				if !flag.Changed ||
+				_, acceptsMultiple := flag.Value.(SliceValue)
+				acceptsMultiple = acceptsMultiple ||
 					strings.Contains(flag.Value.Type(), "Slice") ||
 					strings.Contains(flag.Value.Type(), "Array") ||
-					strings.HasPrefix(flag.Value.Type(), "stringTo") {
+					strings.HasPrefix(flag.Value.Type(), "stringTo")
+
+				if !flag.Changed || acceptsMultiple {
 					// If the flag is not already present, or if it can be specified multiple times (Array, Slice, or stringTo)
 					// we suggest it as a completion
 					completions = append(completions, getFlagNameCompletions(flag, toComplete)...)
diff --git a/completions_test.go b/completions_test.go
index df153fcf..a8f378eb 100644
--- a/completions_test.go
+++ b/completions_test.go
@@ -671,6 +671,29 @@ func TestFlagNameCompletionInGoWithDesc(t *testing.T) {
 	}
 }
 
+// customMultiString is a custom Value type that accepts multiple values,
+// but does not include "Slice" or "Array" in its "Type" string.
+type customMultiString []string
+
+var _ SliceValue = (*customMultiString)(nil)
+
+func (s *customMultiString) String() string {
+	return fmt.Sprintf("%v", *s)
+}
+
+func (s *customMultiString) Set(v string) error {
+	*s = append(*s, v)
+	return nil
+}
+
+func (s *customMultiString) Type() string {
+	return "multi string"
+}
+
+func (s *customMultiString) GetSlice() []string {
+	return *s
+}
+
 func TestFlagNameCompletionRepeat(t *testing.T) {
 	rootCmd := &Command{
 		Use: "root",
@@ -693,6 +716,8 @@ func TestFlagNameCompletionRepeat(t *testing.T) {
 	sliceFlag := rootCmd.Flags().Lookup("slice")
 	rootCmd.Flags().BoolSliceP("bslice", "b", nil, "bool slice flag")
 	bsliceFlag := rootCmd.Flags().Lookup("bslice")
+	rootCmd.Flags().VarP(&customMultiString{}, "multi", "m", "multi string flag")
+	multiFlag := rootCmd.Flags().Lookup("multi")
 
 	// Test that flag names are not repeated unless they are an array or slice
 	output, err := executeCommand(rootCmd, ShellCompNoDescRequestCmd, "--first", "1", "--")
@@ -706,6 +731,7 @@ func TestFlagNameCompletionRepeat(t *testing.T) {
 		"--array",
 		"--bslice",
 		"--help",
+		"--multi",
 		"--second",
 		"--slice",
 		":4",
@@ -728,6 +754,7 @@ func TestFlagNameCompletionRepeat(t *testing.T) {
 		"--array",
 		"--bslice",
 		"--help",
+		"--multi",
 		"--slice",
 		":4",
 		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
@@ -737,7 +764,7 @@ func TestFlagNameCompletionRepeat(t *testing.T) {
 	}
 
 	// Test that flag names are not repeated unless they are an array or slice
-	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "--slice", "1", "--slice=2", "--array", "val", "--bslice", "true", "--")
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "--slice", "1", "--slice=2", "--array", "val", "--bslice", "true", "--multi", "val", "--")
 	if err != nil {
 		t.Errorf("Unexpected error: %v", err)
 	}
@@ -745,12 +772,14 @@ func TestFlagNameCompletionRepeat(t *testing.T) {
 	sliceFlag.Changed = false
 	arrayFlag.Changed = false
 	bsliceFlag.Changed = false
+	multiFlag.Changed = false
 
 	expected = strings.Join([]string{
 		"--array",
 		"--bslice",
 		"--first",
 		"--help",
+		"--multi",
 		"--second",
 		"--slice",
 		":4",
@@ -768,6 +797,7 @@ func TestFlagNameCompletionRepeat(t *testing.T) {
 	// Reset the flag for the next command
 	sliceFlag.Changed = false
 	arrayFlag.Changed = false
+	multiFlag.Changed = false
 
 	expected = strings.Join([]string{
 		"--array",
@@ -778,6 +808,8 @@ func TestFlagNameCompletionRepeat(t *testing.T) {
 		"-f",
 		"--help",
 		"-h",
+		"--multi",
+		"-m",
 		"--second",
 		"-s",
 		"--slice",
@@ -797,6 +829,7 @@ func TestFlagNameCompletionRepeat(t *testing.T) {
 	// Reset the flag for the next command
 	sliceFlag.Changed = false
 	arrayFlag.Changed = false
+	multiFlag.Changed = false
 
 	expected = strings.Join([]string{
 		"-a",