First try at better zsh completions:

A very basic POC. Need to refactor to generate completion
structure before passing to the template to avoid repeated
computations.

What works:
  * Real zsh completion (not built on bash)
  * Basic flags (with long flag and optional shorthand)
  * Basic filename completion indication (not with file extensions though)

What's missing:
  * File extensions to filename completions
  * Positional args
  * Do we require handling only short flags?

Gopkg.lock

@@ -0,0 +1,132 @@
+# This file is autogenerated, do not edit; changes may be undone by the next 'dep ensure'.
+
+
+[[projects]]
+  name = "github.com/cpuguy83/go-md2man"
+  packages = ["md2man"]
+  revision = "20f5889cbdc3c73dbd2862796665e7c465ade7d1"
+  version = "v1.0.8"
+
+[[projects]]
+  name = "github.com/fsnotify/fsnotify"
+  packages = ["."]
+  revision = "c2828203cd70a50dcccfb2761f8b1f8ceef9a8e9"
+  version = "v1.4.7"
+
+[[projects]]
+  branch = "master"
+  name = "github.com/hashicorp/hcl"
+  packages = [
+    ".",
+    "hcl/ast",
+    "hcl/parser",
+    "hcl/printer",
+    "hcl/scanner",
+    "hcl/strconv",
+    "hcl/token",
+    "json/parser",
+    "json/scanner",
+    "json/token"
+  ]
+  revision = "23c074d0eceb2b8a5bfdbb271ab780cde70f05a8"
+
+[[projects]]
+  name = "github.com/inconshreveable/mousetrap"
+  packages = ["."]
+  revision = "76626ae9c91c4f2a10f34cad8ce83ea42c93bb75"
+  version = "v1.0"
+
+[[projects]]
+  name = "github.com/magiconair/properties"
+  packages = ["."]
+  revision = "c3beff4c2358b44d0493c7dda585e7db7ff28ae6"
+  version = "v1.7.6"
+
+[[projects]]
+  branch = "master"
+  name = "github.com/mitchellh/go-homedir"
+  packages = ["."]
+  revision = "b8bc1bf767474819792c23f32d8286a45736f1c6"
+
+[[projects]]
+  branch = "master"
+  name = "github.com/mitchellh/mapstructure"
+  packages = ["."]
+  revision = "a4e142e9c047c904fa2f1e144d9a84e6133024bc"
+
+[[projects]]
+  name = "github.com/pelletier/go-toml"
+  packages = ["."]
+  revision = "acdc4509485b587f5e675510c4f2c63e90ff68a8"
+  version = "v1.1.0"
+
+[[projects]]
+  name = "github.com/russross/blackfriday"
+  packages = ["."]
+  revision = "4048872b16cc0fc2c5fd9eacf0ed2c2fedaa0c8c"
+  version = "v1.5"
+
+[[projects]]
+  name = "github.com/spf13/afero"
+  packages = [
+    ".",
+    "mem"
+  ]
+  revision = "bb8f1927f2a9d3ab41c9340aa034f6b803f4359c"
+  version = "v1.0.2"
+
+[[projects]]
+  name = "github.com/spf13/cast"
+  packages = ["."]
+  revision = "8965335b8c7107321228e3e3702cab9832751bac"
+  version = "v1.2.0"
+
+[[projects]]
+  branch = "master"
+  name = "github.com/spf13/jwalterweatherman"
+  packages = ["."]
+  revision = "7c0cea34c8ece3fbeb2b27ab9b59511d360fb394"
+
+[[projects]]
+  branch = "master"
+  name = "github.com/spf13/pflag"
+  packages = ["."]
+  revision = "6a877ebacf28c5fc79846f4fcd380a5d9872b997"
+
+[[projects]]
+  branch = "master"
+  name = "github.com/spf13/viper"
+  packages = ["."]
+  revision = "aafc9e6bc7b7bb53ddaa75a5ef49a17d6e654be5"
+
+[[projects]]
+  branch = "master"
+  name = "golang.org/x/sys"
+  packages = ["unix"]
+  revision = "37707fdb30a5b38865cfb95e5aab41707daec7fd"
+
+[[projects]]
+  branch = "master"
+  name = "golang.org/x/text"
+  packages = [
+    "internal/gen",
+    "internal/triegen",
+    "internal/ucd",
+    "transform",
+    "unicode/cldr",
+    "unicode/norm"
+  ]
+  revision = "4e4a3210bb54bb31f6ab2cdca2edcc0b50c420c1"
+
+[[projects]]
+  branch = "v2"
+  name = "gopkg.in/yaml.v2"
+  packages = ["."]
+  revision = "d670f9405373e636a5a2765eea47fac0c9bc91a4"
+
+[solve-meta]
+  analyzer-name = "dep"
+  analyzer-version = 1
+  inputs-digest = "d7aecc8ee6a8b343ebf8c2539348464f2566a18dfc511cd374b2cd76bebb1c6d"
+  solver-name = "gps-cdcl"
+  solver-version = 1

Gopkg.toml

@@ -0,0 +1,54 @@
+# Gopkg.toml example
+#
+# Refer to https://github.com/golang/dep/blob/master/docs/Gopkg.toml.md
+# for detailed Gopkg.toml documentation.
+#
+# required = ["github.com/user/thing/cmd/thing"]
+# ignored = ["github.com/user/project/pkgX", "bitbucket.org/user/project/pkgA/pkgY"]
+#
+# [[constraint]]
+#   name = "github.com/user/project"
+#   version = "1.0.0"
+#
+# [[constraint]]
+#   name = "github.com/user/project2"
+#   branch = "dev"
+#   source = "github.com/myfork/project2"
+#
+# [[override]]
+#   name = "github.com/x/y"
+#   version = "2.4.0"
+#
+# [prune]
+#   non-go = false
+#   go-tests = true
+#   unused-packages = true
+
+
+[[constraint]]
+  name = "github.com/cpuguy83/go-md2man"
+  version = "1.0.8"
+
+[[constraint]]
+  name = "github.com/inconshreveable/mousetrap"
+  version = "1.0.0"
+
+[[constraint]]
+  branch = "master"
+  name = "github.com/mitchellh/go-homedir"
+
+[[constraint]]
+  name = "github.com/spf13/pflag"
+  branch = "master"
+
+[[constraint]]
+  name = "github.com/spf13/viper"
+  branch = "master"
+
+[[constraint]]
+  branch = "v2"
+  name = "gopkg.in/yaml.v2"
+
+[prune]
+  go-tests = true
+  unused-packages = true

zsh_completions.go

@@ -1,11 +1,82 @@
 package cobra
 
 import (
-	"bytes"
 	"fmt"
 	"io"
 	"os"
 	"strings"
+	"text/template"
+
+	"github.com/spf13/pflag"
+)
+
+var (
+	funcMap = template.FuncMap{
+		"constructPath": constructPath,
+		"subCmdList":    subCmdList,
+		"extractFlags":  extractFlags,
+		"simpleFlag":    simpleFlag,
+	}
+	zshCompletionText = `
+{{/* for pflag.Flag (specifically annotations) */}}
+{{define "flagAnnotations" -}}
+{{with index .Annotations "cobra_annotation_bash_completion_filename_extensions"}}:filename:_files{{end}}
+{{- end}}
+
+{{/* for pflag.Flag with short and long options */}}
+{{define "complexFlag" -}}
+"(-{{.Shorthand}} --{{.Name}})"{-{{.Shorthand}},--{{.Name}}}"[{{.Usage}}]{{template "flagAnnotations" .}}"
+{{- end}}
+
+{{/* for pflag.Flag with either short or long options */}}
+{{define "simpleFlag" -}}
+"{{with .Name}}--{{.}}{{else}}-{{.Shorthand}}{{end}}[{{.Usage}}]{{template "flagAnnotations" .}}"
+{{- end}}
+
+{{/* should accept Command (that contains subcommands) as parameter */}}
+{{define "argumentsC" -}}
+function {{constructPath .}} {
+  local line
+
+  _arguments -C \
+{{range extractFlags . -}}
+{{"    "}}{{if simpleFlag .}}{{template "simpleFlag" .}}{{else}}{{template "complexFlag" .}}{{end}} \
+{{end}}    "1: :({{subCmdList .}})" \
+    "*::arg:->args"
+
+    case $line[1] in {{- range .Commands}}
+        {{.Use}})
+            {{constructPath .}}
+            ;;
+{{end}}    esac
+}
+{{range .Commands}}
+{{template "selectCmdTemplate" .}}
+{{- end}}
+{{- end}}
+
+{{/* should accept Command without subcommands as parameter */}}
+{{define "arguments" -}}
+function {{constructPath .}} {
+{{with extractFlags . -}}
+{{ "  _arguments" -}}
+{{range .}} \
+    {{if simpleFlag .}}{{template "simpleFlag" .}}{{else}}{{template "complexFlag" .}}{{end -}}
+{{end}}
+{{end -}}
+}
+{{- end}}
+
+{{define "selectCmdTemplate" -}}
+{{if .Commands}}{{template "argumentsC" .}}{{else}}{{template "arguments" .}}{{end}}
+{{- end}}
+
+{{define "Main" -}}
+#compdef _{{.Use}} {{.Use}}
+
+{{template "selectCmdTemplate" .}}
+{{end}}
+`
 )
 
 // GenZshCompletionFile generates zsh completion file.
@@ -21,106 +92,56 @@ func (c *Command) GenZshCompletionFile(filename string) error {
 
 // GenZshCompletion generates a zsh completion file and writes to the passed writer.
 func (c *Command) GenZshCompletion(w io.Writer) error {
-	buf := new(bytes.Buffer)
+	tmpl, err := template.New("Main").Funcs(funcMap).Parse(zshCompletionText)
-
+	if err != nil {
-	writeHeader(buf, c)
+		return fmt.Errorf("error creating zsh completion template: %v", err)
-	maxDepth := maxDepth(c)
+	}
-	writeLevelMapping(buf, maxDepth)
+	return tmpl.Execute(w, c)
-	writeLevelCases(buf, maxDepth, c)
-
-	_, err := buf.WriteTo(w)
-	return err
 }
 
-func writeHeader(w io.Writer, cmd *Command) {
+func constructPath(c *Command) string {
-	fmt.Fprintf(w, "#compdef %s\n\n", cmd.Name())
+	var path []string
+	tmpCmd := c
+	path = append(path, tmpCmd.Use)
+
+	for {
+		if !tmpCmd.HasParent() {
+			break
+		}
+		tmpCmd = tmpCmd.Parent()
+		path = append(path, tmpCmd.Use)
+	}
+
+	// reverse path
+	for left, right := 0, len(path)-1; left < right; left, right = left+1, right-1 {
+		path[left], path[right] = path[right], path[left]
+	}
+
+	return "_" + strings.Join(path, "_")
 }
 
-func maxDepth(c *Command) int {
+// subCmdList returns a space separated list of subcommands names
-	if len(c.Commands()) == 0 {
+func subCmdList(c *Command) string {
-		return 0
+	var subCmds []string
+
+	for _, cmd := range c.Commands() {
+		subCmds = append(subCmds, cmd.Use)
 	}
-	maxDepthSub := 0
+
-	for _, s := range c.Commands() {
+	return strings.Join(subCmds, " ")
-		subDepth := maxDepth(s)
-		if subDepth > maxDepthSub {
-			maxDepthSub = subDepth
-		}
-	}
-	return 1 + maxDepthSub
 }
 
-func writeLevelMapping(w io.Writer, numLevels int) {
+func extractFlags(c *Command) []*pflag.Flag {
-	fmt.Fprintln(w, `_arguments \`)
+	var flags []*pflag.Flag
-	for i := 1; i <= numLevels; i++ {
+	c.LocalFlags().VisitAll(func(f *pflag.Flag) {
-		fmt.Fprintf(w, `  '%d: :->level%d' \`, i, i)
+		flags = append(flags, f)
-		fmt.Fprintln(w)
+	})
-	}
+	c.InheritedFlags().VisitAll(func(f *pflag.Flag) {
-	fmt.Fprintf(w, `  '%d: :%s'`, numLevels+1, "_files")
+		flags = append(flags, f)
-	fmt.Fprintln(w)
+	})
+	return flags
 }
 
-func writeLevelCases(w io.Writer, maxDepth int, root *Command) {
+func simpleFlag(p *pflag.Flag) bool {
-	fmt.Fprintln(w, "case $state in")
+	return p.Name == "" || p.Shorthand == ""
-	defer fmt.Fprintln(w, "esac")
-
-	for i := 1; i <= maxDepth; i++ {
-		fmt.Fprintf(w, "  level%d)\n", i)
-		writeLevel(w, root, i)
-		fmt.Fprintln(w, "  ;;")
-	}
-	fmt.Fprintln(w, "  *)")
-	fmt.Fprintln(w, "    _arguments '*: :_files'")
-	fmt.Fprintln(w, "  ;;")
-}
-
-func writeLevel(w io.Writer, root *Command, i int) {
-	fmt.Fprintf(w, "    case $words[%d] in\n", i)
-	defer fmt.Fprintln(w, "    esac")
-
-	commands := filterByLevel(root, i)
-	byParent := groupByParent(commands)
-
-	for p, c := range byParent {
-		names := names(c)
-		fmt.Fprintf(w, "      %s)\n", p)
-		fmt.Fprintf(w, "        _arguments '%d: :(%s)'\n", i, strings.Join(names, " "))
-		fmt.Fprintln(w, "      ;;")
-	}
-	fmt.Fprintln(w, "      *)")
-	fmt.Fprintln(w, "        _arguments '*: :_files'")
-	fmt.Fprintln(w, "      ;;")
-
-}
-
-func filterByLevel(c *Command, l int) []*Command {
-	cs := make([]*Command, 0)
-	if l == 0 {
-		cs = append(cs, c)
-		return cs
-	}
-	for _, s := range c.Commands() {
-		cs = append(cs, filterByLevel(s, l-1)...)
-	}
-	return cs
-}
-
-func groupByParent(commands []*Command) map[string][]*Command {
-	m := make(map[string][]*Command)
-	for _, c := range commands {
-		parent := c.Parent()
-		if parent == nil {
-			continue
-		}
-		m[parent.Name()] = append(m[parent.Name()], c)
-	}
-	return m
-}
-
-func names(commands []*Command) []string {
-	ns := make([]string, len(commands))
-	for i, c := range commands {
-		ns[i] = c.Name()
-	}
-	return ns
 }

zsh_completions_test.go

@@ -2,74 +2,92 @@ package cobra
 
 import (
 	"bytes"
-	"strings"
+	"regexp"
 	"testing"
 )
 
-func TestZshCompletion(t *testing.T) {
+func TestGenZshCompletion(t *testing.T) {
+	var debug bool
+	var option string
+
 	tcs := []struct {
 		name                string
 		root                *Command
 		expectedExpressions []string
 	}{
 		{
-			name:                "trivial",
+			name: "simple command",
-			root:                &Command{Use: "trivialapp"},
-			expectedExpressions: []string{"#compdef trivial"},
-		},
-		{
-			name: "linear",
 			root: func() *Command {
-				r := &Command{Use: "linear"}
+				r := &Command{
-
+					Use:  "mycommand",
-				sub1 := &Command{Use: "sub1"}
+					Long: "My Command long description",
-				r.AddCommand(sub1)
+				}
-
+				r.Flags().BoolVar(&debug, "debug", debug, "description")
-				sub2 := &Command{Use: "sub2"}
-				sub1.AddCommand(sub2)
-
-				sub3 := &Command{Use: "sub3"}
-				sub2.AddCommand(sub3)
 				return r
 			}(),
-			expectedExpressions: []string{"sub1", "sub2", "sub3"},
+			expectedExpressions: []string{
+				`function _mycommand {\s+_arguments \\\s+"--debug\[description\]"\s+}`,
+				"#compdef _mycommand mycommand",
+			},
 		},
 		{
-			name: "flat",
+			name: "flags with both long and short flags",
 			root: func() *Command {
-				r := &Command{Use: "flat"}
+				r := &Command{
-				r.AddCommand(&Command{Use: "c1"})
+					Use:  "testcmd",
-				r.AddCommand(&Command{Use: "c2"})
+					Long: "long description",
+				}
+				r.Flags().BoolVarP(&debug, "debug", "d", debug, "debug description")
 				return r
 			}(),
-			expectedExpressions: []string{"(c1 c2)"},
+			expectedExpressions: []string{
+				`"\(-d --debug\)"{-d,--debug}"\[debug description\]"`,
+			},
 		},
 		{
-			name: "tree",
+			name: "command with subcommands",
 			root: func() *Command {
-				r := &Command{Use: "tree"}
+				r := &Command{
-
+					Use:  "rootcmd",
-				sub1 := &Command{Use: "sub1"}
+					Long: "Long rootcmd description",
-				r.AddCommand(sub1)
+				}
-
+				d := &Command{
-				sub11 := &Command{Use: "sub11"}
+					Use:   "subcmd1",
-				sub12 := &Command{Use: "sub12"}
+					Short: "Subcmd1 short descrition",
-
+				}
-				sub1.AddCommand(sub11)
+				e := &Command{
-				sub1.AddCommand(sub12)
+					Use:  "subcmd2",
-
+					Long: "Subcmd2 short description",
-				sub2 := &Command{Use: "sub2"}
+				}
-				r.AddCommand(sub2)
+				r.PersistentFlags().BoolVar(&debug, "debug", debug, "description")
-
+				d.Flags().StringVarP(&option, "option", "o", option, "option description")
-				sub21 := &Command{Use: "sub21"}
+				r.AddCommand(d, e)
-				sub22 := &Command{Use: "sub22"}
-
-				sub2.AddCommand(sub21)
-				sub2.AddCommand(sub22)
-
 				return r
 			}(),
-			expectedExpressions: []string{"(sub11 sub12)", "(sub21 sub22)"},
+			expectedExpressions: []string{
+				`\\\n\s+"1: :\(subcmd1 subcmd2\)" \\\n`,
+				`_arguments \\\n.*"--debug\[description]"`,
+				`_arguments -C \\\n.*"--debug\[description]"`,
+				`function _rootcmd_subcmd1 {`,
+				`function _rootcmd_subcmd1 {`,
+				`_arguments \\\n.*"\(-o --option\)"{-o,--option}"\[option description]" \\\n`,
+			},
+		},
+		{
+			name: "filename completion",
+			root: func() *Command {
+				var file string
+				r := &Command{
+					Use:   "mycmd",
+					Short: "my command short description",
+				}
+				r.Flags().StringVarP(&file, "config", "c", file, "config file")
+				r.MarkFlagFilename("config", "ext")
+				return r
+			}(),
+			expectedExpressions: []string{
+				`\n +"\(-c --config\)"{-c,--config}"\[config file]:filename:_files"`,
+			},
 		},
 	}
 
@@ -77,13 +95,66 @@ func TestZshCompletion(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			buf := new(bytes.Buffer)
 			tc.root.GenZshCompletion(buf)
-			output := buf.String()
+			output := buf.Bytes()
 
-			for _, expectedExpression := range tc.expectedExpressions {
+			for _, expr := range tc.expectedExpressions {
-				if !strings.Contains(output, expectedExpression) {
+				rgx, err := regexp.Compile(expr)
-					t.Errorf("Expected completion to contain %q somewhere; got %q", expectedExpression, output)
+				if err != nil {
+					t.Errorf("error compiling expression (%s): %v", expr, err)
+				}
+				if !rgx.Match(output) {
+					t.Errorf("expeced completion (%s) to match '%s'", buf.String(), expr)
 				}
 			}
 		})
 	}
+
+}
+
+func BenchmarkConstructPath(b *testing.B) {
+	c := &Command{
+		Use:   "main",
+		Long:  "main long description which is very long",
+		Short: "main short description",
+	}
+	d := &Command{
+		Use: "hello",
+	}
+	e := &Command{
+		Use: "world",
+	}
+	c.AddCommand(d)
+	d.AddCommand(e)
+	for i := 0; i < b.N; i++ {
+		res := constructPath(e)
+		if res != "_main_hello_world" {
+			b.Errorf("expeced path to be '_main_hello_world', got %s", res)
+		}
+	}
+}
+
+func TestExtractFlags(t *testing.T) {
+	var debug, cmdc, cmdd bool
+	c := &Command{
+		Use:  "cmdC",
+		Long: "Command C",
+	}
+	c.PersistentFlags().BoolVarP(&debug, "debug", "d", debug, "debug mode")
+	c.Flags().BoolVar(&cmdc, "cmd-c", cmdc, "Command C")
+	d := &Command{
+		Use:  "CmdD",
+		Long: "Command D",
+	}
+	d.Flags().BoolVar(&cmdd, "cmd-d", cmdd, "Command D")
+	c.AddCommand(d)
+
+	resC := extractFlags(c)
+	resD := extractFlags(d)
+
+	if len(resC) != 2 {
+		t.Errorf("expected Command C to return 2 flags, got %d", len(resC))
+	}
+	if len(resD) != 2 {
+		t.Errorf("expected Command D to return 2 flags, got %d", len(resD))
+	}
 }

zsh_template.tmpl

@@ -0,0 +1,56 @@
+{{define "complexFlag"}}{{ /* for pflag.Flag with short and long options */ -}}
+"(-{{.Shorthand}} --{{.Name}})"\{-{{.Shorthand}}, --{{.Name}}\}[{{.Usage}}]
+{{- end}}
+
+{{define "simpleFlag"}}{{ /* for pflag.Flag with either short or long options */ -}}
+"{{with .Name}}-{{.}}{{else}}--{{.Shorthand}}{{end}}[{{.Usage}}]"
+{{- end}}
+
+{{define "argumentsC"}}
+{{- /* should accept Command (that contains subcommands) as parameter */ -}}
+function {{constructPath .}} {
+  local line
+
+  _arguments -C \
+{{range extractFlags . -}}
+    {{if simpleFlag .}}{{template "simpleFlag" .}}{{else}}{{template "complexFlag" .}}{{end}} \
+{{end -}}
+    "1: :({{subCmdList .}})" \
+    "*:arg:->args"
+
+    case $line[1] in
+{{range .Commands -}}
+        {{.Use}})
+            {{constructPath .}}
+            ;;
+{{end -}}
+    esac
+}
+{{range .Commands -}}
+
+{{template "selectCmdTemplate" .}}
+{{end -}}
+{{end}}
+
+{{define "arguments"}}
+function {{constructPath .}} {
+{{- /* should accept Command without subcommands as parameter */ -}}
+{{with extractFlags . -}}
+  _arguments \
+{{range .}}
+    {{if simpleFlag .}}{{template "simpleFlag" .}}{{else}}{{template "complexFlag" .}}{{end}} \
+{{end -}}
+{{ /* leave this line empty because of the last backslash */ }}
+{{end -}}
+}
+{{end}}
+
+{{define "selectCmdTemplate" -}}
+{{with .Commands}}{{template "argumentsC" .}}{{else}}{{template "arguments" .}}{{end}}
+{{end -}}
+
+{{define "Main"}}
+#compdef _{{.Use}} {{.Use}}
+
+{{template "selectCmdTemplate" .}}
+{{end}}