diff --git a/zsh_completions.go b/zsh_completions.go index 889c22e2..6f3b4b45 100644 --- a/zsh_completions.go +++ b/zsh_completions.go @@ -6,6 +6,8 @@ import ( "io" "os" "strings" + + flag "github.com/spf13/pflag" ) // GenZshCompletionFile generates zsh completion file. @@ -19,14 +21,23 @@ func (c *Command) GenZshCompletionFile(filename string) error { return c.GenZshCompletion(outFile) } +func argName(cmd *Command) string { + for cmd.HasParent() { + cmd = cmd.Parent() + } + name := fmt.Sprintf("%s_cmd_args", cmd.Name()) + return strings.Replace(name, "-", "_",-1) +} // GenZshCompletion generates a zsh completion file and writes to the passed writer. func (c *Command) GenZshCompletion(w io.Writer) error { buf := new(bytes.Buffer) writeHeader(buf, c) maxDepth := maxDepth(c) - writeLevelMapping(buf, maxDepth) + fmt.Fprintf(buf, "_%s() {\n", c.Name()) + writeLevelMapping(buf, maxDepth, c) writeLevelCases(buf, maxDepth, c) + fmt.Fprintf(buf, "}\n_%s \"$@\"\n", c.Name()) _, err := buf.WriteTo(w) return err @@ -50,77 +61,102 @@ func maxDepth(c *Command) int { return 1 + maxDepthSub } -func writeLevelMapping(w io.Writer, numLevels int) { - fmt.Fprintln(w, `_arguments \`) +func writeLevelMapping(w io.Writer, numLevels int, root *Command) { + fmt.Fprintln(w, `local context curcontext="$curcontext" state line`) + fmt.Fprintln(w, `typeset -A opt_args`) + fmt.Fprintln(w, `_arguments -C \`) for i := 1; i <= numLevels; i++ { - fmt.Fprintf(w, ` '%d: :->level%d' \`, i, i) - fmt.Fprintln(w) + fmt.Fprintf(w, " '%d: :->level%d' \\\n", i, i) } - fmt.Fprintf(w, ` '%d: :%s'`, numLevels+1, "_files") - fmt.Fprintln(w) + fmt.Fprintf(w, " $%s \\\n", argName(root)) + fmt.Fprintln(w, ` '*: :_files'`) } func writeLevelCases(w io.Writer, maxDepth int, root *Command) { fmt.Fprintln(w, "case $state in") - 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, " ;;") + fmt.Fprintln(w, "esac") } -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, " ;;") +func writeLevel(w io.Writer, root *Command, l int) { + fmt.Fprintf(w, " level%d)\n", l) + fmt.Fprintf(w, " case $words[%d] in\n", l) + for _, c := range filterByLevel(root, l) { + writeCommandArgsBlock(w, c) } fmt.Fprintln(w, " *)") fmt.Fprintln(w, " _arguments '*: :_files'") fmt.Fprintln(w, " ;;") + fmt.Fprintln(w, " esac") + fmt.Fprintln(w, " ;;") +} +func writeCommandArgsBlock(w io.Writer, c *Command) { + names := commandNames(c) + flags := commandFlags(c) + if len(names) > 0 || len(flags) > 0 { + fmt.Fprintf(w, " %s)\n", c.Name()) + defer fmt.Fprintln(w, " ;;") + } + if len(flags) > 0 { + fmt.Fprintf(w, " %s=(\n", argName(c)) + for _, flag := range flags { + fmt.Fprintf(w, " %s\n", flag) + } + fmt.Fprintln(w, " )") + } + if len(names) > 0 { + fmt.Fprintf(w, " _values 'command' '%s'\n", strings.Join(names, "' '")) + } } 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 + commands := []*Command{c} + for i := 1; i < l; i++ { + var nextLevel []*Command + for _, c := range commands { + if c.HasSubCommands() { + nextLevel = append(nextLevel, c.Commands()...) + } } - m[parent.Name()] = append(m[parent.Name()], c) + commands = nextLevel } - return m + + return commands } -func names(commands []*Command) []string { +func commandNames(command *Command) []string { + commands := command.Commands() ns := make([]string, len(commands)) for i, c := range commands { - ns[i] = c.Name() + commandMsg := c.Name() + if len(c.Short) > 0 { + commandMsg += fmt.Sprintf("[%s]", c.Short) + } + ns[i] = commandMsg } return ns } + +func commandFlags(command *Command) []string { + flags := command.Flags() + ns := make([]string, 0) + flags.VisitAll(func(flag *flag.Flag) { + var flagMsg string + if len(flag.Shorthand) > 0 { + flagMsg = fmt.Sprintf("{-%s,--%s}", flag.Shorthand, flag.Name) + } else { + flagMsg = fmt.Sprintf("--%s", flag.Name) + } + if len(flag.Usage) > 0 { + flagMsg += fmt.Sprintf("'[%s]'", flag.Usage) + } + ns = append(ns, flagMsg) + }) + return ns +} diff --git a/zsh_completions_test.go b/zsh_completions_test.go index 34e69496..f8d5388a 100644 --- a/zsh_completions_test.go +++ b/zsh_completions_test.go @@ -42,7 +42,7 @@ func TestZshCompletion(t *testing.T) { r.AddCommand(&Command{Use: "c2"}) return r }(), - expectedExpressions: []string{"(c1 c2)"}, + expectedExpressions: []string{"'c1' 'c2'"}, }, { name: "tree", @@ -69,7 +69,7 @@ func TestZshCompletion(t *testing.T) { return r }(), - expectedExpressions: []string{"(sub11 sub12)", "(sub21 sub22)"}, + expectedExpressions: []string{"'sub11' 'sub12'", "'sub21' 'sub22'"}, }, }