zsh-completions: implemented argument completion.

This commit is contained in:
Haim Ashkenazi 2018-03-23 13:09:56 +03:00 committed by Steve Francia
parent d262154093
commit edbb6712e2
3 changed files with 283 additions and 11 deletions

View File

@ -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()

View File

@ -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).

View File

@ -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)