fix RegisterFlagCompletionFunc concurrent map writes error (#1423)

* fix-RegisterFlagCompletionFunc-concurrent
* set to root command
* move to non-public fields
This commit is contained in:
silenceshell 2021-07-01 05:49:30 +08:00 committed by GitHub
parent 2dea4f2ef3
commit 3c8a19ecd3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 12 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

@ -142,6 +142,9 @@ type Command struct {
// that we can use on every pflag set and children commands // that we can use on every pflag set and children commands
globNormFunc func(f *flag.FlagSet, name string) flag.NormalizedName globNormFunc func(f *flag.FlagSet, name string) flag.NormalizedName
//flagCompletionFunctions is map of flag completion functions.
flagCompletionFunctions map[*flag.Flag]func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective)
// usageFunc is usage func defined by user. // usageFunc is usage func defined by user.
usageFunc func(*Command) error usageFunc func(*Command) error
// usageTemplate is usage template defined by user. // usageTemplate is usage template defined by user.

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
@ -94,10 +91,15 @@ 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 {
root := c.Root()
if _, exists := 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 root.flagCompletionFunctions == nil {
root.flagCompletionFunctions = map[*pflag.Flag]func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective){}
}
root.flagCompletionFunctions[flag] = f
return nil return nil
} }
@ -374,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
} }