From 3c8a19ecd384dc4229545bb174310a50c493f4ae Mon Sep 17 00:00:00 2001 From: silenceshell Date: Thu, 1 Jul 2021 05:49:30 +0800 Subject: [PATCH] fix RegisterFlagCompletionFunc concurrent map writes error (#1423) * fix-RegisterFlagCompletionFunc-concurrent * set to root command * move to non-public fields --- bash_completions.go | 2 +- command.go | 3 +++ completions.go | 14 ++++++++------ 3 files changed, 12 insertions(+), 7 deletions(-) diff --git a/bash_completions.go b/bash_completions.go index b47f898a..925e6e78 100644 --- a/bash_completions.go +++ b/bash_completions.go @@ -512,7 +512,7 @@ func writeLocalNonPersistentFlag(buf io.StringWriter, flag *pflag.Flag) { // Setup annotations for go completions for registered flags 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 // 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 diff --git a/command.go b/command.go index 2cc18891..9f33a461 100644 --- a/command.go +++ b/command.go @@ -142,6 +142,9 @@ type Command struct { // that we can use on every pflag set and children commands 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 func(*Command) error // usageTemplate is usage template defined by user. diff --git a/completions.go b/completions.go index 5dfea616..5f925ccf 100644 --- a/completions.go +++ b/completions.go @@ -17,9 +17,6 @@ const ( 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 // can be instructed to have once completions have been provided. type ShellCompDirective int @@ -94,10 +91,15 @@ func (c *Command) RegisterFlagCompletionFunc(flagName string, f func(cmd *Comman if flag == nil { 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) } - 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 } @@ -374,7 +376,7 @@ func (c *Command) getCompletions(args []string) (*Command, []string, ShellCompDi // Find the completion function for the flag or command var completionFn func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective) if flag != nil { - completionFn = flagCompletionFunctions[flag] + completionFn = c.Root().flagCompletionFunctions[flag] } else { completionFn = finalCmd.ValidArgsFunction }