diff --git a/zsh_completions.go b/zsh_completions.go index 490eb02..68bb5c6 100644 --- a/zsh_completions.go +++ b/zsh_completions.go @@ -1,20 +1,29 @@ package cobra import ( + "encoding/json" "fmt" "io" "os" + "sort" "strings" "text/template" "github.com/spf13/pflag" ) +const ( + zshCompArgumentAnnotation = "cobra_annotations_zsh_completion_argument_annotation" + zshCompArgumentFilenameComp = "cobra_annotations_zsh_completion_argument_file_completion" + zshCompArgumentWordComp = "cobra_annotations_zsh_completion_argument_word_completion" +) + var ( zshCompFuncMap = template.FuncMap{ "genZshFuncName": zshCompGenFuncName, "extractFlags": zshCompExtractFlag, "genFlagEntryForZshArguments": zshCompGenFlagEntryForArguments, + "extractArgsCompletions": zshCompExtractArgumentCompletionHintsForRendering, } zshCompletionText = ` {{/* should accept Command (that contains subcommands) as parameter */}} @@ -53,7 +62,8 @@ function {{$cmdPath}} { function {{genZshFuncName .}} { {{" _arguments"}}{{range extractFlags .}} \ {{genFlagEntryForZshArguments . -}} -{{end}} +{{end}}{{range extractArgsCompletions .}} \ + {{.}}{{end}} } {{end}} @@ -73,6 +83,19 @@ function {{genZshFuncName .}} { ` ) +// zshCompArgsAnnotation is used to encode/decode zsh completion for +// arguments to/from Command.Annotations. +type zshCompArgsAnnotation map[int]zshCompArgHint + +type zshCompArgHint struct { + // Indicates the type of the completion to use. One of: + // zshCompArgumentFilenameComp or zshCompArgumentWordComp + Tipe string `json:"type"` + + // A value for the type above (globs for file completion or words) + Options []string `json:"options"` +} + // GenZshCompletionFile generates zsh completion file. func (c *Command) GenZshCompletionFile(filename string) error { outFile, err := os.Create(filename) @@ -95,6 +118,130 @@ func (c *Command) GenZshCompletion(w io.Writer) error { return tmpl.Execute(w, c.Root()) } +// MarkZshCompPositionalArgumentFile marks the specified argument (first +// argument is 1) as completed by file selection. patterns (e.g. "*.txt") are +// optional - if not provided the completion will search for all files. +func (c *Command) MarkZshCompPositionalArgumentFile(argPosition int, patterns ...string) error { + if argPosition < 1 { + return fmt.Errorf("Invalid argument position (%d)", argPosition) + } + annotation, err := c.zshCompGetArgsAnnotations() + if err != nil { + return err + } + if c.zshcompArgsAnnotationnIsDuplicatePosition(annotation, argPosition) { + return fmt.Errorf("Duplicate annotation for positional argument at index %d", argPosition) + } + annotation[argPosition] = zshCompArgHint{ + Tipe: zshCompArgumentFilenameComp, + Options: patterns, + } + return c.zshCompSetArgsAnnotations(annotation) +} + +// MarkZshCompPositionalArgumentWords marks the specified positional argument +// (first argument is 1) as completed by the provided words. At east one word +// must be provided, spaces within words will be offered completion with +// "word\ word". +func (c *Command) MarkZshCompPositionalArgumentWords(argPosition int, words ...string) error { + if argPosition < 1 { + return fmt.Errorf("Invalid argument position (%d)", argPosition) + } + if len(words) == 0 { + return fmt.Errorf("Trying to set empty word list for positional argument %d", argPosition) + } + annotation, err := c.zshCompGetArgsAnnotations() + if err != nil { + return err + } + if c.zshcompArgsAnnotationnIsDuplicatePosition(annotation, argPosition) { + return fmt.Errorf("Duplicate annotation for positional argument at index %d", argPosition) + } + annotation[argPosition] = zshCompArgHint{ + Tipe: zshCompArgumentWordComp, + Options: words, + } + return c.zshCompSetArgsAnnotations(annotation) +} + +func zshCompExtractArgumentCompletionHintsForRendering(c *Command) ([]string, error) { + var result []string + annotation, err := c.zshCompGetArgsAnnotations() + if err != nil { + return nil, err + } + for k, v := range annotation { + s, err := zshCompRenderZshCompArgHint(k, v) + if err != nil { + return nil, err + } + result = append(result, s) + } + if len(c.ValidArgs) > 0 { + if _, positionOneExists := annotation[1]; !positionOneExists { + s, err := zshCompRenderZshCompArgHint(1, zshCompArgHint{ + Tipe: zshCompArgumentWordComp, + Options: c.ValidArgs, + }) + if err != nil { + return nil, err + } + result = append(result, s) + } + } + sort.Strings(result) + return result, nil +} + +func zshCompRenderZshCompArgHint(i int, z zshCompArgHint) (string, error) { + switch t := z.Tipe; t { + case zshCompArgumentFilenameComp: + var globs []string + for _, g := range z.Options { + globs = append(globs, fmt.Sprintf(`-g "%s"`, g)) + } + return fmt.Sprintf(`'%d: :_files %s'`, i, strings.Join(globs, " ")), nil + case zshCompArgumentWordComp: + var words []string + for _, w := range z.Options { + words = append(words, fmt.Sprintf("%q", w)) + } + return fmt.Sprintf(`'%d: :(%s)'`, i, strings.Join(words, " ")), nil + default: + return "", fmt.Errorf("Invalid zsh argument completion annotation: %s", t) + } +} + +func (c *Command) zshcompArgsAnnotationnIsDuplicatePosition(annotation zshCompArgsAnnotation, position int) bool { + _, dup := annotation[position] + return dup +} + +func (c *Command) zshCompGetArgsAnnotations() (zshCompArgsAnnotation, error) { + annotation := make(zshCompArgsAnnotation) + annotationString, ok := c.Annotations[zshCompArgumentAnnotation] + if !ok { + return annotation, nil + } + err := json.Unmarshal([]byte(annotationString), &annotation) + if err != nil { + return annotation, fmt.Errorf("Error unmarshaling zsh argument annotation: %v", err) + } + return annotation, nil +} + +func (c *Command) zshCompSetArgsAnnotations(annotation zshCompArgsAnnotation) error { + jsn, err := json.Marshal(annotation) + if err != nil { + return fmt.Errorf("Error marshaling zsh argument annotation: %v", err) + } + if c.Annotations == nil { + c.Annotations = make(map[string]string) + } + c.Annotations[zshCompArgumentAnnotation] = string(jsn) + return nil +} + func zshCompGenFuncName(c *Command) string { if c.HasParent() { return zshCompGenFuncName(c.Parent()) + "_" + c.Name() diff --git a/zsh_completions.md b/zsh_completions.md index c218179..95242d3 100644 --- a/zsh_completions.md +++ b/zsh_completions.md @@ -14,10 +14,25 @@ The generated completion script should be put somewhere in your `$fpath` named flag value - if it's empty then completion will expect an argument. * Flags of one of the various `*Arrary` and `*Slice` types supports multiple specifications (with or without argument depending on the specific type). +* Completion of positional arguments using the following rules: + * Argument position for all options below starts at `1`. If argument position + `0` is requested it will raise an error. + * Use `command.MarkZshCompPositionalArgumentFile` to complete filenames. Glob + patterns (e.g. `"*.log"`) are optional - if not specified it will offer to + complete all file types. + * Use `command.MarkZshCompPositionalArgumentWords` to offer specific words for + completion. At least one word is required. + * It's possible to specify completion for some arguments and leave some + unspecified (e.g. offer words for second argument but nothing for first + argument). This will cause no completion for first argument but words + completion for second argument. + * If no argument completion was specified for 1st argument (but optionally was + specified for 2nd) and the command has `ValidArgs` it will be used as + completion options for 1st argument. + * Argument completions only offered for commands with no subcommands. ### What's not yet Supported -* Positional argument completion are not supported yet. * Custom completion scripts are not supported yet (We should probably create zsh specific one, doesn't make sense to re-use the bash one as the functions will be different). diff --git a/zsh_completions_test.go b/zsh_completions_test.go index 4ef2e2f..976cbfc 100644 --- a/zsh_completions_test.go +++ b/zsh_completions_test.go @@ -58,7 +58,7 @@ func TestGenZshCompletion(t *testing.T) { } d := &Command{ Use: "subcmd1", - Short: "Subcmd1 short descrition", + Short: "Subcmd1 short description", Run: emptyRun, } e := &Command{ @@ -135,7 +135,7 @@ func TestGenZshCompletion(t *testing.T) { skip: "--version and --help are currently not generated when not running on root command", }, { - name: "zsh generation should run on root commannd", + name: "zsh generation should run on root command", root: func() *Command { r := genTestCommand("root", false) s := genTestCommand("sub1", true) @@ -157,6 +157,63 @@ func TestGenZshCompletion(t *testing.T) { `--private\[Don'\\''t show public info]`, }, }, + { + name: "argument completion for file with and without patterns", + root: func() *Command { + r := genTestCommand("root", true) + r.MarkZshCompPositionalArgumentFile(1, "*.log") + r.MarkZshCompPositionalArgumentFile(2) + return r + }(), + expectedExpressions: []string{ + `'1: :_files -g "\*.log"' \\\n\s+'2: :_files`, + }, + }, + { + name: "argument zsh completion for words", + root: func() *Command { + r := genTestCommand("root", true) + r.MarkZshCompPositionalArgumentWords(1, "word1", "word2") + return r + }(), + expectedExpressions: []string{ + `'1: :\("word1" "word2"\)`, + }, + }, + { + name: "argument completion for words with spaces", + root: func() *Command { + r := genTestCommand("root", true) + r.MarkZshCompPositionalArgumentWords(1, "single", "multiple words") + return r + }(), + expectedExpressions: []string{ + `'1: :\("single" "multiple words"\)'`, + }, + }, + { + name: "argument completion when command has ValidArgs and no annotation for argument completion", + root: func() *Command { + r := genTestCommand("root", true) + r.ValidArgs = []string{"word1", "word2"} + return r + }(), + expectedExpressions: []string{ + `'1: :\("word1" "word2"\)'`, + }, + }, + { + name: "argument completion when command has ValidArgs and no annotation for argument at argPosition 1", + root: func() *Command { + r := genTestCommand("root", true) + r.ValidArgs = []string{"word1", "word2"} + r.MarkZshCompPositionalArgumentFile(2) + return r + }(), + expectedExpressions: []string{ + `'1: :\("word1" "word2"\)' \\`, + }, + }, } for _, tc := range tcs { @@ -178,7 +235,7 @@ func TestGenZshCompletion(t *testing.T) { t.Errorf("error compiling expression (%s): %v", expr, err) } if !rgx.Match(output) { - t.Errorf("expeced completion (%s) to match '%s'", buf.String(), expr) + t.Errorf("expected completion (%s) to match '%s'", buf.String(), expr) } } }) @@ -192,7 +249,7 @@ func TestGenZshCompletionHidden(t *testing.T) { expectedExpressions []string }{ { - name: "hidden commmands", + name: "hidden commands", root: func() *Command { r := &Command{ Use: "main", @@ -255,8 +312,61 @@ func TestGenZshCompletionHidden(t *testing.T) { } } +func TestMarkZshCompPositionalArgumentFile(t *testing.T) { + t.Run("Doesn't allow overwriting existing positional argument", func(t *testing.T) { + c := &Command{} + if err := c.MarkZshCompPositionalArgumentFile(1, "*.log"); err != nil { + t.Errorf("Received error when we shouldn't have: %v\n", err) + } + if err := c.MarkZshCompPositionalArgumentFile(1); err == nil { + t.Error("Didn't receive an error when trying to overwrite argument position") + } + }) + + t.Run("Refuses to accept argPosition less then 1", func(t *testing.T) { + c := &Command{} + err := c.MarkZshCompPositionalArgumentFile(0, "*") + if err == nil { + t.Fatal("Error was not thrown when indicating argument position 0") + } + if !strings.Contains(err.Error(), "position") { + t.Errorf("expected error message '%s' to contain 'position'", err.Error()) + } + }) +} + +func TestMarkZshCompPositionalArgumentWords(t *testing.T) { + t.Run("Doesn't allow overwriting existing positional argument", func(t *testing.T) { + c := &Command{} + if err := c.MarkZshCompPositionalArgumentFile(1, "*.log"); err != nil { + t.Errorf("Received error when we shouldn't have: %v\n", err) + } + if err := c.MarkZshCompPositionalArgumentWords(1, "hello"); err == nil { + t.Error("Didn't receive an error when trying to overwrite argument position") + } + }) + + t.Run("Doesn't allow calling without words", func(t *testing.T) { + c := &Command{} + if err := c.MarkZshCompPositionalArgumentWords(0); err == nil { + t.Error("Should not allow saving empty word list for annotation") + } + }) + + t.Run("Refuses to accept argPosition less then 1", func(t *testing.T) { + c := &Command{} + err := c.MarkZshCompPositionalArgumentWords(0, "word") + if err == nil { + t.Fatal("Should not allow setting argument position less then 1") + } + if !strings.Contains(err.Error(), "position") { + t.Errorf("Expected error '%s' to contain 'position' but didn't", err.Error()) + } + }) +} + func BenchmarkMediumSizeConstruct(b *testing.B) { - root := constructLargeCommandHeirarchy() + root := constructLargeCommandHierarchy() // if err := root.GenZshCompletionFile("_mycmd"); err != nil { // b.Error(err) // } @@ -296,7 +406,7 @@ func TestExtractFlags(t *testing.T) { } } -func constructLargeCommandHeirarchy() *Command { +func constructLargeCommandHierarchy() *Command { var config, st1, st2 string var long, debug bool var in1, in2 int @@ -308,7 +418,7 @@ func constructLargeCommandHeirarchy() *Command { panic(err) } s1 := genTestCommand("sub1", true) - s1.Flags().BoolVar(&long, "long", long, "long descriptin") + s1.Flags().BoolVar(&long, "long", long, "long description") s1.Flags().BoolSliceVar(&verbose, "verbose", verbose, "verbose description") s1.Flags().StringArray("option", []string{}, "various options") s2 := genTestCommand("sub2", true) @@ -320,8 +430,8 @@ func constructLargeCommandHeirarchy() *Command { s1_1.Flags().StringVar(&st2, "st2", st2, "st2 description") s1_2 := genTestCommand("sub1sub2", true) s1_3 := genTestCommand("sub1sub3", true) - s1_3.Flags().IntVar(&in1, "int1", in1, "int1 descriptionn") - s1_3.Flags().IntVar(&in2, "int2", in2, "int2 descriptionn") + s1_3.Flags().IntVar(&in1, "int1", in1, "int1 description") + s1_3.Flags().IntVar(&in2, "int2", in2, "int2 description") s1_3.Flags().StringArrayP("option", "O", []string{}, "more options") s2_1 := genTestCommand("sub2sub1", true) s2_2 := genTestCommand("sub2sub2", true)