zsh: added custom completions

command.go

@@ -68,6 +68,9 @@ type Command struct {
 	// BashCompletionFunction is custom functions used by the bash autocompletion generator.
 	BashCompletionFunction string
 
+	// ZshCompletionFunction is custom functions used by the zsh autocompletion generator.
+	ZshCompletionFunction string
+
 	// Deprecated defines, if this command is deprecated and should print this string when used.
 	Deprecated string
 

shell_completions.go

@@ -34,6 +34,12 @@ func (c *Command) MarkFlagCustom(name string, f string) error {
 	return MarkFlagCustom(c.Flags(), name, f)
 }
 
+// MarkPFlagCustom adds the BashCompCustom annotation to the named pflag, if it exists.
+// Generated shell autocompletion will call the function f for the flag.
+func (c *Command) MarkPFlagCustom(name string, f string) error {
+	return MarkFlagCustom(c.PersistentFlags(), name, f)
+}
+
 // MarkPersistentFlagFilename instructs the various shell completion
 // implementations to limit completions for this persistent flag to the
 // specified extensions (patterns).

zsh_completions.go

@@ -14,6 +14,7 @@ import (
 
 const (
 	zshCompArgumentAnnotation   = "cobra_annotations_zsh_completion_argument_annotation"
+	zshCompArgumentCustomComp   = "cobra_annotations_zsh_completion_argument_custom_completion"
 	zshCompArgumentFilenameComp = "cobra_annotations_zsh_completion_argument_file_completion"
 	zshCompArgumentWordComp     = "cobra_annotations_zsh_completion_argument_word_completion"
 	zshCompDirname              = "cobra_annotations_zsh_dirname"
@@ -79,6 +80,8 @@ function {{genZshFuncName .}} {
 {{define "Main" -}}
 #compdef _{{.Name}} {{.Name}}
 
+{{.ZshCompletionFunction}}
+
 {{template "selectCmdTemplate" .}}
 {{end}}
 `
@@ -119,6 +122,26 @@ func (c *Command) GenZshCompletion(w io.Writer) error {
 	return tmpl.Execute(w, c.Root())
 }
 
+// MarkZshCompPositionalArgumentCustom marks the specified argument (first
+// argument is 1) as custom.
+func (c *Command) MarkZshCompPositionalArgumentCustom(argPosition int, function 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:    zshCompArgumentCustomComp,
+		Options: []string{function},
+	}
+	return c.zshCompSetArgsAnnotations(annotation)
+}
+
 // 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.
@@ -208,6 +231,8 @@ func zshCompRenderZshCompArgHint(i int, z zshCompArgHint) (string, error) {
 			words = append(words, fmt.Sprintf("%q", w))
 		}
 		return fmt.Sprintf(`'%d: :(%s)'`, i, strings.Join(words, " ")), nil
+	case zshCompArgumentCustomComp:
+		return fmt.Sprintf(`'%d: :%s'`, i, z.Options[0]), nil
 	default:
 		return "", fmt.Errorf("Invalid zsh argument completion annotation: %s", t)
 	}
@@ -310,7 +335,7 @@ func zshCompGenFlagEntryExtras(f *pflag.Flag) string {
 		return ""
 	}
 
-	extras := ":" // allow options for flag (even without assistance)
+	extras := ":()" // allow options for flag (even without assistance)
 	for key, values := range f.Annotations {
 		switch key {
 		case zshCompDirname:
@@ -320,6 +345,8 @@ func zshCompGenFlagEntryExtras(f *pflag.Flag) string {
 			for _, pattern := range values {
 				extras = extras + fmt.Sprintf(` -g "%s"`, pattern)
 			}
+		case BashCompCustom:
+			extras = ": :" + values[0]
 		}
 	}
 

zsh_completions_test.go

@@ -77,7 +77,7 @@ func TestGenZshCompletion(t *testing.T) {
 				`_arguments -C \\\n.*'--debug\[description]'`,
 				`function _rootcmd_subcmd1 {`,
 				`function _rootcmd_subcmd1 {`,
-				`_arguments \\\n.*'\(-o --option\)'{-o,--option}'\[option description]:' \\\n`,
+				`_arguments \\\n.*'\(-o --option\)'{-o,--option}'\[option description]:\(\)' \\\n`,
 			},
 		},
 		{
@@ -379,6 +379,32 @@ func TestMarkZshCompPositionalArgumentWords(t *testing.T) {
 		}
 	})
 }
+func TestMarkZshCompPositionalArgumentCustom(t *testing.T) {
+	t.Run("testPositionalCustom", func(t *testing.T) {
+		c := &Command{}
+		c.ZshCompletionFunction = `
+function __custom_function {
+  _values 'test' a b c
+}`
+		c.MarkZshCompPositionalArgumentCustom(1, "__custom_function")
+
+		buf := new(bytes.Buffer)
+		err := c.GenZshCompletion(buf)
+		if err != nil {
+			t.Error(err)
+		}
+
+		output := buf.String()
+
+		if !strings.Contains(output, "'1: :__custom_function'") {
+			t.Error("should contain custom function argument")
+		}
+
+		if !strings.Contains(output, "function __custom_function {") {
+			t.Error("should contain custom function")
+		}
+	})
+}
 
 func BenchmarkMediumSizeConstruct(b *testing.B) {
 	root := constructLargeCommandHierarchy()