zsh-completions: implemented argument completion.

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

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

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(&in1, "int1", in1, "int1 description")
-	s1_3.Flags().IntVar(&in2, "int2", in2, "int2 descriptionn")
+	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)