Don't use global map for flag completions.

This commit is contained in:
maxlandon 2023-08-08 04:14:47 +02:00
parent fd865a44e3
commit f1f260eb59
No known key found for this signature in database
GPG key ID: 2DE5C14975A86900
3 changed files with 23 additions and 16 deletions

View file

@ -534,9 +534,9 @@ func writeLocalNonPersistentFlag(buf io.StringWriter, flag *pflag.Flag) {
// prepareCustomAnnotationsForFlags setup annotations for go completions for registered flags // prepareCustomAnnotationsForFlags setup annotations for go completions for registered flags
func prepareCustomAnnotationsForFlags(cmd *Command) { func prepareCustomAnnotationsForFlags(cmd *Command) {
flagCompletionMutex.RLock() cmd.flagCompletionMutex.RLock()
defer flagCompletionMutex.RUnlock() defer cmd.flagCompletionMutex.RUnlock()
for flag := range flagCompletionFunctions { for flag := range cmd.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
@ -644,6 +644,7 @@ func writeCmdAliases(buf io.StringWriter, cmd *Command) {
WriteStringAndCheck(buf, ` fi`) WriteStringAndCheck(buf, ` fi`)
WriteStringAndCheck(buf, "\n") WriteStringAndCheck(buf, "\n")
} }
func writeArgAliases(buf io.StringWriter, cmd *Command) { func writeArgAliases(buf io.StringWriter, cmd *Command) {
WriteStringAndCheck(buf, " noun_aliases=()\n") WriteStringAndCheck(buf, " noun_aliases=()\n")
sort.Strings(cmd.ArgAliases) sort.Strings(cmd.ArgAliases)

View file

@ -26,6 +26,7 @@ import (
"path/filepath" "path/filepath"
"sort" "sort"
"strings" "strings"
"sync"
flag "github.com/spf13/pflag" flag "github.com/spf13/pflag"
) )
@ -158,6 +159,13 @@ 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
// flagsCompletions contrains completions for arbitrary lists of flags.
// Those flags may or may not actually strictly belong to the command in the function,
// but registering completions for them through the command allows for garbage-collecting.
flagCompletionFunctions map[*flag.Flag]func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective)
// lock for reading and writing from flagCompletionFunctions
flagCompletionMutex *sync.RWMutex
// 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

@ -18,7 +18,6 @@ import (
"fmt" "fmt"
"os" "os"
"strings" "strings"
"sync"
"github.com/spf13/pflag" "github.com/spf13/pflag"
) )
@ -33,10 +32,10 @@ const (
) )
// Global map of flag completion functions. Make sure to use flagCompletionMutex before you try to read and write from it. // Global map of flag completion functions. Make sure to use flagCompletionMutex before you try to read and write from it.
var flagCompletionFunctions = map[*pflag.Flag]func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective){} // var flagCompletionFunctions = map[*pflag.Flag]func(cmd *Command, args []string, toComplete string) ([]string, ShellCompDirective){}
//
// lock for reading and writing from flagCompletionFunctions // // lock for reading and writing from flagCompletionFunctions
var flagCompletionMutex = &sync.RWMutex{} // var flagCompletionMutex = &sync.RWMutex{}
// 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.
@ -135,13 +134,13 @@ 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)
} }
flagCompletionMutex.Lock() c.flagCompletionMutex.Lock()
defer flagCompletionMutex.Unlock() defer c.flagCompletionMutex.Unlock()
if _, exists := flagCompletionFunctions[flag]; exists { if _, exists := c.flagCompletionFunctions[flag]; exists {
return fmt.Errorf("RegisterFlagCompletionFunc: flag '%s' already registered", flagName) return fmt.Errorf("RegisterFlagCompletionFunc: flag '%s' already registered", flagName)
} }
flagCompletionFunctions[flag] = f c.flagCompletionFunctions[flag] = f
return nil return nil
} }
@ -479,9 +478,9 @@ 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 && flagCompletion { if flag != nil && flagCompletion {
flagCompletionMutex.RLock() finalCmd.flagCompletionMutex.RLock()
completionFn = flagCompletionFunctions[flag] completionFn = finalCmd.flagCompletionFunctions[flag]
flagCompletionMutex.RUnlock() finalCmd.flagCompletionMutex.RUnlock()
} else { } else {
completionFn = finalCmd.ValidArgsFunction completionFn = finalCmd.ValidArgsFunction
} }
@ -805,7 +804,6 @@ to your powershell profile.
return cmd.Root().GenPowerShellCompletion(out) return cmd.Root().GenPowerShellCompletion(out)
} }
return cmd.Root().GenPowerShellCompletionWithDesc(out) return cmd.Root().GenPowerShellCompletionWithDesc(out)
}, },
} }
if haveNoDescFlag { if haveNoDescFlag {