cobra/zsh_completions.go

package cobra

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"strings"
)

// GenZshCompletionFile generates zsh completion file.
func (c *Command) GenZshCompletionFile(filename string) error {
	outFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer outFile.Close()

	return c.GenZshCompletion(outFile)
}

// 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)
	writeLevelCases(buf, maxDepth, c)

	_, err := buf.WriteTo(w)
	return err
}

func writeHeader(w io.Writer, cmd *Command) {
	fmt.Fprintf(w, "#compdef %s\n\n", cmd.Name())
}

func maxDepth(c *Command) int {
	if len(c.Commands()) == 0 {
		return 0
	}
	maxDepthSub := 0
	for _, s := range c.Commands() {
		subDepth := maxDepth(s)
		if subDepth > maxDepthSub {
			maxDepthSub = subDepth
		}
	}
	return 1 + maxDepthSub
}

func writeLevelMapping(w io.Writer, numLevels int) {
	fmt.Fprintln(w, `_arguments \`)
	for i := 1; i <= numLevels; i++ {
		fmt.Fprintf(w, `  '%d: :->level%d' \`, i, i)
		fmt.Fprintln(w)
	}
	fmt.Fprintf(w, `  '%d: :%s'`, numLevels+1, "_files")
	fmt.Fprintln(w)
}

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, "  ;;")
}

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
}