From b755f40ce77306c134af523e3a8756126d0b1b75 Mon Sep 17 00:00:00 2001
From: tim <timaz98@ya.ru>
Date: Fri, 17 Jun 2022 12:23:23 +0500
Subject: [PATCH] feat(completion): support no-space completions for short
 flags

---
 completions.go      |  19 ++++++
 completions_test.go | 137 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 156 insertions(+)

diff --git a/completions.go b/completions.go
index 3de4e87a..832fbfc1 100644
--- a/completions.go
+++ b/completions.go
@@ -543,6 +543,25 @@ func checkIfFlagCompletion(finalCmd *Command, args []string, lastArg string) (*p
 			lastArg = lastArg[index+1:]
 			flagWithEqual = true
 		} else {
+
+			// Find first shorthand non-boolean flag
+			// Following characters are treated as an argument
+			// e.g. `-ini` => n is a flag, 'i' - argument to complete
+			// https://github.com/spf13/cobra/issues/1629
+			if len(lastArg) > 1 && !strings.HasPrefix(lastArg, "--") {
+				i := 1
+				for ; i < len(lastArg)-1; i++ {
+					flagName = lastArg[i : i+1]
+					f := findFlag(finalCmd, flagName)
+					if f != nil && len(f.NoOptDefVal) == 0 {
+						flagName = lastArg[i : i+1]
+						lastArg = lastArg[i+1:]
+
+						return f, args, lastArg, nil
+					}
+				}
+			}
+
 			// Normal flag completion
 			return nil, args, lastArg, nil
 		}
diff --git a/completions_test.go b/completions_test.go
index 28a3228e..a9b2f97d 100644
--- a/completions_test.go
+++ b/completions_test.go
@@ -2691,3 +2691,140 @@ func TestFixedCompletions(t *testing.T) {
 		t.Errorf("expected: %q, got: %q", expected, output)
 	}
 }
+
+func TestCompleteWithShortFlagNoSpace(t *testing.T) {
+	rootCmd := &Command{
+		Use: "root",
+		Run: emptyRun,
+	}
+
+	f := rootCmd.Flags()
+	f.BoolP("interactive", "i", false, "short flag 1")
+	f.BoolP("timeout", "t", false, "short flag 2")
+	f.StringP("namespace", "n", "defValue", "file namespace")
+	_ = rootCmd.RegisterFlagCompletionFunc("short3", func(*Command, []string, string) ([]string, ShellCompDirective) {
+		return []string{"works"}, ShellCompDirectiveNoFileComp
+	})
+
+	namespaceComps := []string{"infra-test", "infra-prod", "prod"}
+
+	_ = rootCmd.RegisterFlagCompletionFunc("namespace", func(c *Command, args []string, toComplete string) ([]string, ShellCompDirective) {
+		comps := make([]string, 0, len(namespaceComps))
+		for _, comp := range namespaceComps {
+			if strings.HasPrefix(comp, toComplete) {
+				comps = append(comps, comp)
+			}
+		}
+		if len(comps) == 0 {
+			comps = namespaceComps
+		}
+		return comps, ShellCompDirectiveNoFileComp
+	})
+
+	// Test that shorthand flag works
+	output, err := executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-n")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected := strings.Join([]string{
+		"-n",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+
+	// Test that boolean + string shorthand flags work
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-in")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected = strings.Join([]string{
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+
+	// Test that boolean + string with single character value shorthand flags work
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-ini")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected = strings.Join([]string{
+		"infra-test",
+		"infra-prod",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+
+	// Test that boolean + string with single character value shorthand flags work
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-inp")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected = strings.Join([]string{
+		"prod",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+
+	// Test that multiple boolean + string with single character value shorthand flags work
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-itni")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected = strings.Join([]string{
+		"infra-test",
+		"infra-prod",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+
+	// Test that boolean + string with multiple character value shorthand flags work
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-ininfra")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected = strings.Join([]string{
+		"infra-test",
+		"infra-prod",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+
+	// Test that boolean + string with multiple character value shorthand flags work
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "-ininfra-t")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected = strings.Join([]string{
+		"infra-test",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+}