fix-RegisterFlagCompletionFunc-concurrent

This commit is contained in:
silenceshell 2021-06-24 11:40:44 +08:00
parent 9a432671fd
commit 4ee33b3c5b
3 changed files with 11 additions and 7 deletions

View file

@ -512,7 +512,7 @@ func writeLocalNonPersistentFlag(buf io.StringWriter, flag *pflag.Flag) {
// Setup annotations for go completions for registered flags // Setup annotations for go completions for registered flags
func prepareCustomAnnotationsForFlags(cmd *Command) { func prepareCustomAnnotationsForFlags(cmd *Command) {
for flag := range flagCompletionFunctions { for flag := range cmd.Root().flagCompletionFunctions {
// Make sure the completion script calls the __*_go_custom_completion function for // Make sure the completion script calls the __*_go_custom_completion function for
// every registered flag. We need to do this here (and not when the flag was registered // every registered flag. We need to do this here (and not when the flag was registered
// for completion) so that we can know the root command name for the prefix // for completion) so that we can know the root command name for the prefix

View file

@ -88,6 +88,9 @@ type Command struct {
// group commands. // group commands.
Annotations map[string]string Annotations map[string]string
//flagCompletionFunctions is map of flag completion functions.
flagCompletionFunctions map[*flag.Flag]func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective)
// Version defines the version for this command. If this value is non-empty and the command does not // Version defines the version for this command. If this value is non-empty and the command does not
// define a "version" flag, a "version" boolean flag will be added to the command and, if specified, // define a "version" flag, a "version" boolean flag will be added to the command and, if specified,
// will print content of the "Version" variable. A shorthand "v" flag will also be added if the // will print content of the "Version" variable. A shorthand "v" flag will also be added if the

View file

@ -17,9 +17,6 @@ const (
ShellCompNoDescRequestCmd = "__completeNoDesc" ShellCompNoDescRequestCmd = "__completeNoDesc"
) )
// Global map of flag completion functions.
var flagCompletionFunctions = map[*pflag.Flag]func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective){}
// ShellCompDirective is a bit map representing the different behaviors the shell // ShellCompDirective is a bit map representing the different behaviors the shell
// can be instructed to have once completions have been provided. // can be instructed to have once completions have been provided.
type ShellCompDirective int type ShellCompDirective int
@ -95,10 +92,14 @@ func (c *Command) RegisterFlagCompletionFunc(flagName string, f func(cmd *Comman
if flag == nil { if flag == nil {
return fmt.Errorf("RegisterFlagCompletionFunc: flag '%s' does not exist", flagName) return fmt.Errorf("RegisterFlagCompletionFunc: flag '%s' does not exist", flagName)
} }
if _, exists := flagCompletionFunctions[flag]; exists {
if _, exists := c.Root().flagCompletionFunctions[flag]; exists {
return fmt.Errorf("RegisterFlagCompletionFunc: flag '%s' already registered", flagName) return fmt.Errorf("RegisterFlagCompletionFunc: flag '%s' already registered", flagName)
} }
flagCompletionFunctions[flag] = f if c.Root().flagCompletionFunctions == nil {
c.Root().flagCompletionFunctions = map[*pflag.Flag]func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective){}
}
c.Root().flagCompletionFunctions[flag] = f
return nil return nil
} }
@ -375,7 +376,7 @@ func (c *Command) getCompletions(args []string) (*Command, []string, ShellCompDi
// Find the completion function for the flag or command // Find the completion function for the flag or command
var completionFn func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective) var completionFn func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective)
if flag != nil { if flag != nil {
completionFn = flagCompletionFunctions[flag] completionFn = c.Root().flagCompletionFunctions[flag]
} else { } else {
completionFn = finalCmd.ValidArgsFunction completionFn = finalCmd.ValidArgsFunction
} }