Go spf13/Cobra custom flag types
Cobra allows you to define custom value types to be used as flags through the pflag.(*FlagSet).Var()
method (from the https://github.com/spf13/pflag package, that is used by Cobra). You have to make a new type that implements the pflag.Value
interface:
type Value interface {
String() string
Set(string) error
Type() string
}
Example type definition:
type myEnum string
const (
myEnumFoo myEnum = "foo"
myEnumBar myEnum = "bar"
myEnumMoo myEnum = "moo"
)
// String is used both by fmt.Print and by Cobra in help text
func (e *myEnum) String() string {
return string(*e)
}
// Set must have pointer receiver so it doesn't change the value of a copy
func (e *myEnum) Set(v string) error {
switch v {
case "foo", "bar", "moo":
*e = myEnum(v)
return nil
default:
return errors.New(`must be one of "foo", "bar", or "moo"`)
}
}
// Type is only used in help text
func (e *myEnum) Type() string {
return "myEnum"
}
Example registration:
func init() {
var flagMyEnum = myEnumFoo
var myCmd = &cobra.Command{
Use: "mycmd",
Short: "A brief description of your command",
Run: func(cmd *cobra.Command, args []string) {
fmt.Println("myenum value:", flagMyEnum)
},
}
rootCmd.AddCommand(myCmd)
myCmd.Flags().Var(&flagMyEnum, "myenum", `my custom enum. allowed: "foo", "bar", "moo"`)
}
Example usage: (notice the first line in the console output below)
$ go run . mycmd --myenum raz
Error: invalid argument "raz" for "--myenum" flag: must be one of "foo", "bar", or "moo"
Usage:
main mycmd [flags]
Flags:
-h, --help help for mycmd
--myenum myEnum my custom enum. allowed: "foo", "bar", "moo" (default foo)
exit status 1
$ go run . mycmd --myenum bar
myenum value: bar
Completions
To add autocompletion to this, the somewhat hidden documentation page cobra/shell_completions.md#completions-for-flags is at great assistance. For our example, you would add something like this:
func init() {
// ...
myCmd.Flags().Var(&flagMyEnum, "myenum", `my custom enum. allowed: "foo", "bar", "moo"`)
myCmd.RegisterFlagCompletionFunc("myenum", myEnumCompletion)
}
// myEnumCompletion should probably live next to the myEnum definition
func myEnumCompletion(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
return []string{
"foo\thelp text for foo",
"bar\thelp text for bar",
"moo\thelp text for moo",
}, cobra.ShellCompDirectiveDefault
}
Example usage:
$ go build -o main .
$ source <(main completion bash)
$ main mycmd --myenum <TAB><TAB>
bar (help text for bar)
foo (help text for foo)
moo (help text for moo)