From c79b2dd3dcff10e5c1b4e23f98e737edcb336752 Mon Sep 17 00:00:00 2001
From: Marc Khouzam <marc.khouzam@montreal.ca>
Date: Tue, 14 Dec 2021 07:16:21 -0500
Subject: [PATCH] Remove __complete cmd for program without subcmds

Fixes #1562

Programs that don't have sub-commands can accept any number of args.
However, when doing shell completion for such programs, within the
__complete code this very __complete command makes it that the program
suddenly has a sub-command, and the call to Find() -> legacyArgs() will
then return an error if there are more than one argument on the
command-line being completed.

To avoid this, we first remove the __complete command in such a case so
as to get back to having no sub-commands.

Signed-off-by: Marc Khouzam <marc.khouzam@montreal.ca>
---
 completions.go      | 12 +++++++++++-
 completions_test.go | 45 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/completions.go b/completions.go
index ac25a9d03..9ecd56a47 100644
--- a/completions.go
+++ b/completions.go
@@ -228,7 +228,17 @@ func (c *Command) getCompletions(args []string) (*Command, []string, ShellCompDi
 	if c.Root().TraverseChildren {
 		finalCmd, finalArgs, err = c.Root().Traverse(trimmedArgs)
 	} else {
-		finalCmd, finalArgs, err = c.Root().Find(trimmedArgs)
+		// For Root commands that don't specify any value for their Args fields, when we call
+		// Find(), if those Root commands don't have any sub-commands, they will accept arguments.
+		// However, because we have added the __complete sub-command in the current code path, the
+		// call to Find() -> legacyArgs() will return an error if there are any arguments.
+		// To avoid this, we first remove the __complete command to get back to having no sub-commands.
+		rootCmd := c.Root()
+		if len(rootCmd.Commands()) == 1 {
+			rootCmd.RemoveCommand(c)
+		}
+
+		finalCmd, finalArgs, err = rootCmd.Find(trimmedArgs)
 	}
 	if err != nil {
 		// Unable to find the real command. E.g., <program> someInvalidCmd <TAB>
diff --git a/completions_test.go b/completions_test.go
index 2d6c11dc3..6a62f0cc8 100644
--- a/completions_test.go
+++ b/completions_test.go
@@ -2619,3 +2619,48 @@ func TestCompleteWithDisableFlagParsing(t *testing.T) {
 		t.Errorf("expected: %q, got: %q", expected, output)
 	}
 }
+
+func TestCompleteWithRootAndLegacyArgs(t *testing.T) {
+	// Test a lonely root command which uses legacyArgs().  In such a case, the root
+	// command should accept any number of arguments and completion should behave accordingly.
+	rootCmd := &Command{
+		Use:  "root",
+		Args: nil, // Args must be nil to trigger the legacyArgs() function
+		Run:  emptyRun,
+		ValidArgsFunction: func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective) {
+			return []string{"arg1", "arg2"}, ShellCompDirectiveNoFileComp
+		},
+	}
+
+	// Make sure the first arg is completed
+	output, err := executeCommand(rootCmd, ShellCompNoDescRequestCmd, "")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected := strings.Join([]string{
+		"arg1",
+		"arg2",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+
+	// Make sure the completion of arguments continues
+	output, err = executeCommand(rootCmd, ShellCompNoDescRequestCmd, "arg1", "")
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	expected = strings.Join([]string{
+		"arg1",
+		"arg2",
+		":4",
+		"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n")
+
+	if output != expected {
+		t.Errorf("expected: %q, got: %q", expected, output)
+	}
+}