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)