Update shell completion to respect flag groups (#1659)

Signed-off-by: Marc Khouzam <marc.khouzam@montreal.ca>

Co-authored-by: Marc Khouzam <marc.khouzam@montreal.ca>

completions.go

@@ -325,6 +325,9 @@ func (c *Command) getCompletions(args []string) (*Command, []string, ShellCompDi
 	var completions []string
 	var directive ShellCompDirective
 
+	// Enforce flag groups before doing flag completions
+	finalCmd.enforceFlagGroupsForCompletion()
+
 	// Note that we want to perform flagname completion even if finalCmd.DisableFlagParsing==true;
 	// doing this allows for completion of persistent flag names even for commands that disable flag parsing.
 	//

completions_test.go

@@ -2691,3 +2691,189 @@ func TestFixedCompletions(t *testing.T) {
 		t.Errorf("expected: %q, got: %q", expected, output)
 	}
 }
+
+func TestCompletionForGroupedFlags(t *testing.T) {
+	getCmd := func() *Command {
+		rootCmd := &Command{
+			Use: "root",
+			Run: emptyRun,
+		}
+		childCmd := &Command{
+			Use: "child",
+			ValidArgsFunction: func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective) {
+				return []string{"subArg"}, ShellCompDirectiveNoFileComp
+			},
+			Run: emptyRun,
+		}
+		rootCmd.AddCommand(childCmd)
+
+		rootCmd.PersistentFlags().Int("ingroup1", -1, "ingroup1")
+		rootCmd.PersistentFlags().String("ingroup2", "", "ingroup2")
+
+		childCmd.Flags().Bool("ingroup3", false, "ingroup3")
+		childCmd.Flags().Bool("nogroup", false, "nogroup")
+
+		// Add flags to a group
+		childCmd.MarkFlagsRequiredTogether("ingroup1", "ingroup2", "ingroup3")
+
+		return rootCmd
+	}
+
+	// Each test case uses a unique command from the function above.
+	testcases := []struct {
+		desc           string
+		args           []string
+		expectedOutput string
+	}{
+		{
+			desc: "flags in group not suggested without - prefix",
+			args: []string{"child", ""},
+			expectedOutput: strings.Join([]string{
+				"subArg",
+				":4",
+				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
+		},
+		{
+			desc: "flags in group suggested with - prefix",
+			args: []string{"child", "-"},
+			expectedOutput: strings.Join([]string{
+				"--ingroup1",
+				"--ingroup2",
+				"--ingroup3",
+				"--nogroup",
+				":4",
+				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
+		},
+		{
+			desc: "when flag in group present, other flags in group suggested even without - prefix",
+			args: []string{"child", "--ingroup2", "value", ""},
+			expectedOutput: strings.Join([]string{
+				"--ingroup1",
+				"--ingroup3",
+				"subArg",
+				":4",
+				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
+		},
+		{
+			desc: "when all flags in group present, flags not suggested without - prefix",
+			args: []string{"child", "--ingroup1", "8", "--ingroup2", "value2", "--ingroup3", ""},
+			expectedOutput: strings.Join([]string{
+				"subArg",
+				":4",
+				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
+		},
+		{
+			desc: "group ignored if some flags not applicable",
+			args: []string{"--ingroup2", "value", ""},
+			expectedOutput: strings.Join([]string{
+				"child",
+				"completion",
+				"help",
+				":4",
+				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.desc, func(t *testing.T) {
+			c := getCmd()
+			args := []string{ShellCompNoDescRequestCmd}
+			args = append(args, tc.args...)
+			output, err := executeCommand(c, args...)
+			switch {
+			case err == nil && output != tc.expectedOutput:
+				t.Errorf("expected: %q, got: %q", tc.expectedOutput, output)
+			case err != nil:
+				t.Errorf("Unexpected error %q", err)
+			}
+		})
+	}
+}
+
+func TestCompletionForMutuallyExclusiveFlags(t *testing.T) {
+	getCmd := func() *Command {
+		rootCmd := &Command{
+			Use: "root",
+			Run: emptyRun,
+		}
+		childCmd := &Command{
+			Use: "child",
+			ValidArgsFunction: func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective) {
+				return []string{"subArg"}, ShellCompDirectiveNoFileComp
+			},
+			Run: emptyRun,
+		}
+		rootCmd.AddCommand(childCmd)
+
+		rootCmd.PersistentFlags().IntSlice("ingroup1", []int{1}, "ingroup1")
+		rootCmd.PersistentFlags().String("ingroup2", "", "ingroup2")
+
+		childCmd.Flags().Bool("ingroup3", false, "ingroup3")
+		childCmd.Flags().Bool("nogroup", false, "nogroup")
+
+		// Add flags to a group
+		childCmd.MarkFlagsMutuallyExclusive("ingroup1", "ingroup2", "ingroup3")
+
+		return rootCmd
+	}
+
+	// Each test case uses a unique command from the function above.
+	testcases := []struct {
+		desc           string
+		args           []string
+		expectedOutput string
+	}{
+		{
+			desc: "flags in mutually exclusive group not suggested without the - prefix",
+			args: []string{"child", ""},
+			expectedOutput: strings.Join([]string{
+				"subArg",
+				":4",
+				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
+		},
+		{
+			desc: "flags in mutually exclusive group suggested with the - prefix",
+			args: []string{"child", "-"},
+			expectedOutput: strings.Join([]string{
+				"--ingroup1",
+				"--ingroup2",
+				"--ingroup3",
+				"--nogroup",
+				":4",
+				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
+		},
+		{
+			desc: "when flag in mutually exclusive group present, other flags in group not suggested even with the - prefix",
+			args: []string{"child", "--ingroup1", "8", "-"},
+			expectedOutput: strings.Join([]string{
+				"--ingroup1", // Should be suggested again since it is a slice
+				"--nogroup",
+				":4",
+				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
+		},
+		{
+			desc: "group ignored if some flags not applicable",
+			args: []string{"--ingroup1", "8", "-"},
+			expectedOutput: strings.Join([]string{
+				"--ingroup1",
+				"--ingroup2",
+				":4",
+				"Completion ended with directive: ShellCompDirectiveNoFileComp", ""}, "\n"),
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.desc, func(t *testing.T) {
+			c := getCmd()
+			args := []string{ShellCompNoDescRequestCmd}
+			args = append(args, tc.args...)
+			output, err := executeCommand(c, args...)
+			switch {
+			case err == nil && output != tc.expectedOutput:
+				t.Errorf("expected: %q, got: %q", tc.expectedOutput, output)
+			case err != nil:
+				t.Errorf("Unexpected error %q", err)
+			}
+		})
+	}
+}

flag_groups.go

@@ -172,3 +172,52 @@ func sortedKeys(m map[string]map[string]bool) []string {
 	sort.Strings(keys)
 	return keys
 }
+
+// enforceFlagGroupsForCompletion will do the following:
+// - when a flag in a group is present, other flags in the group will be marked required
+// - when a flag in a mutually exclusive group is present, other flags in the group will be marked as hidden
+// This allows the standard completion logic to behave appropriately for flag groups
+func (c *Command) enforceFlagGroupsForCompletion() {
+	if c.DisableFlagParsing {
+		return
+	}
+
+	flags := c.Flags()
+	groupStatus := map[string]map[string]bool{}
+	mutuallyExclusiveGroupStatus := map[string]map[string]bool{}
+	c.Flags().VisitAll(func(pflag *flag.Flag) {
+		processFlagForGroupAnnotation(flags, pflag, requiredAsGroup, groupStatus)
+		processFlagForGroupAnnotation(flags, pflag, mutuallyExclusive, mutuallyExclusiveGroupStatus)
+	})
+
+	// If a flag that is part of a group is present, we make all the other flags
+	// of that group required so that the shell completion suggests them automatically
+	for flagList, flagnameAndStatus := range groupStatus {
+		for _, isSet := range flagnameAndStatus {
+			if isSet {
+				// One of the flags of the group is set, mark the other ones as required
+				for _, fName := range strings.Split(flagList, " ") {
+					_ = c.MarkFlagRequired(fName)
+				}
+			}
+		}
+	}
+
+	// If a flag that is mutually exclusive to others is present, we hide the other
+	// flags of that group so the shell completion does not suggest them
+	for flagList, flagnameAndStatus := range mutuallyExclusiveGroupStatus {
+		for flagName, isSet := range flagnameAndStatus {
+			if isSet {
+				// One of the flags of the mutually exclusive group is set, mark the other ones as hidden
+				// Don't mark the flag that is already set as hidden because it may be an
+				// array or slice flag and therefore must continue being suggested
+				for _, fName := range strings.Split(flagList, " ") {
+					if fName != flagName {
+						flag := c.Flags().Lookup(fName)
+						flag.Hidden = true
+					}
+				}
+			}
+		}
+	}
+}