From 2d94892a8bec681d65a129e0a94b2e0d87e52cb9 Mon Sep 17 00:00:00 2001
From: Paul Holzinger <45212748+Luap99@users.noreply.github.com>
Date: Mon, 3 May 2021 18:42:00 +0200
Subject: [PATCH] Custom completion handle multiple shorhand flags together
 (#1258)

Flag definitions like `-asd` are not handled correctly by
the custom completion logic. They should be treated as
multiple flags. For details refer to #1257.

Fixes #1257

Signed-off-by: Paul Holzinger <paul.holzinger@web.de>
---
 completions.go      | 23 +++++++++--
 completions_test.go | 93 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 113 insertions(+), 3 deletions(-)

diff --git a/completions.go b/completions.go
index 28d7dd03..6982d9a8 100644
--- a/completions.go
+++ b/completions.go
@@ -469,7 +469,16 @@ func checkIfFlagCompletion(finalCmd *Command, args []string, lastArg string) (*p
 	if len(lastArg) > 0 && lastArg[0] == '-' {
 		if index := strings.Index(lastArg, "="); index >= 0 {
 			// Flag with an =
-			flagName = strings.TrimLeft(lastArg[:index], "-")
+			if strings.HasPrefix(lastArg[:index], "--") {
+				// Flag has full name
+				flagName = lastArg[2:index]
+			} else {
+				// Flag is shorthand
+				// We have to get the last shorthand flag name
+				// e.g. `-asd` => d to provide the correct completion
+				// https://github.com/spf13/cobra/issues/1257
+				flagName = lastArg[index-1 : index]
+			}
 			lastArg = lastArg[index+1:]
 			flagWithEqual = true
 		} else {
@@ -486,8 +495,16 @@ func checkIfFlagCompletion(finalCmd *Command, args []string, lastArg string) (*p
 				// If the flag contains an = it means it has already been fully processed,
 				// so we don't need to deal with it here.
 				if index := strings.Index(prevArg, "="); index < 0 {
-					flagName = strings.TrimLeft(prevArg, "-")
-
+					if strings.HasPrefix(prevArg, "--") {
+						// Flag has full name
+						flagName = prevArg[2:]
+					} else {
+						// Flag is shorthand
+						// We have to get the last shorthand flag name
+						// e.g. `-asd` => d to provide the correct completion
+						// https://github.com/spf13/cobra/issues/1257
+						flagName = prevArg[len(prevArg)-1:]
+					}
 					// Remove the uncompleted flag or else there could be an error created
 					// for an invalid value for that flag
 					trimmedArgs = args[:len(args)-1]
diff --git a/completions_test.go b/completions_test.go
index 3e16bd04..9bd6a410 100644
--- a/completions_test.go
+++ b/completions_test.go
@@ -2201,3 +2201,96 @@ func TestCompleteCompletion(t *testing.T) {
 		}
 	}
 }
+
+func TestMultipleShorthandFlagCompletion(t *testing.T) {
+	rootCmd := &Command{
+		Use:       "root",
+		ValidArgs: []string{"foo", "bar"},
+		Run:       emptyRun,
+	}
+	f := rootCmd.Flags()
+	f.BoolP("short", "s", false, "short flag 1")
+	f.BoolP("short2", "d", false, "short flag 2")
+	f.StringP("short3", "f", "", "short flag 3")
+	_ = rootCmd.RegisterFlagCompletionFunc("short3", func(*Command, []string, string) ([]string, ShellCompDirective) {
+		return []string{"works"}, ShellCompDirectiveNoFileComp
+	})
+
+	// Test that a single shorthand flag works
+	output, err := executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-s", "")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected := strings.Join([]string{
+		"foo",
+		"bar",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+
+	// Test that multiple boolean shorthand flags work
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-sd", "")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected = strings.Join([]string{
+		"foo",
+		"bar",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+
+	// Test that multiple boolean + string shorthand flags work
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-sdf", "")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected = strings.Join([]string{
+		"works",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+
+	// Test that multiple boolean + string with equal sign shorthand flags work
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-sdf=")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected = strings.Join([]string{
+		"works",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+
+	// Test that multiple boolean + string with equal sign with value shorthand flags work
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-sdf=abc", "")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected = strings.Join([]string{
+		"foo",
+		"bar",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+}