vendor/github.com/spf13/cobra/zsh_completions.go

   1 package cobra
   2
   3 import (
   4         "bytes"
   5         "fmt"
   6         "io"
   7         "os"
   8         "strings"
   9 )
  10
  11 // GenZshCompletionFile generates zsh completion file.
  12 func (c *Command) GenZshCompletionFile(filename string) error {
  13         outFile, err := os.Create(filename)
  14         if err != nil {
  15                 return err
  16         }
  17         defer outFile.Close()
  18
  19         return c.GenZshCompletion(outFile)
  20 }
  21
  22 // GenZshCompletion generates a zsh completion file and writes to the passed writer.
  23 func (c *Command) GenZshCompletion(w io.Writer) error {
  24         buf := new(bytes.Buffer)
  25
  26         writeHeader(buf, c)
  27         maxDepth := maxDepth(c)
  28         writeLevelMapping(buf, maxDepth)
  29         writeLevelCases(buf, maxDepth, c)
  30
  31         _, err := buf.WriteTo(w)
  32         return err
  33 }
  34
  35 func writeHeader(w io.Writer, cmd *Command) {
  36         fmt.Fprintf(w, "#compdef %s\n\n", cmd.Name())
  37 }
  38
  39 func maxDepth(c *Command) int {
  40         if len(c.Commands()) == 0 {
  41                 return 0
  42         }
  43         maxDepthSub := 0
  44         for _, s := range c.Commands() {
  45                 subDepth := maxDepth(s)
  46                 if subDepth > maxDepthSub {
  47                         maxDepthSub = subDepth
  48                 }
  49         }
  50         return 1 + maxDepthSub
  51 }
  52
  53 func writeLevelMapping(w io.Writer, numLevels int) {
  54         fmt.Fprintln(w, `_arguments \`)
  55         for i := 1; i <= numLevels; i++ {
  56                 fmt.Fprintf(w, `  '%d: :->level%d' \`, i, i)
  57                 fmt.Fprintln(w)
  58         }
  59         fmt.Fprintf(w, `  '%d: :%s'`, numLevels+1, "_files")
  60         fmt.Fprintln(w)
  61 }
  62
  63 func writeLevelCases(w io.Writer, maxDepth int, root *Command) {
  64         fmt.Fprintln(w, "case $state in")
  65         defer fmt.Fprintln(w, "esac")
  66
  67         for i := 1; i <= maxDepth; i++ {
  68                 fmt.Fprintf(w, "  level%d)\n", i)
  69                 writeLevel(w, root, i)
  70                 fmt.Fprintln(w, "  ;;")
  71         }
  72         fmt.Fprintln(w, "  *)")
  73         fmt.Fprintln(w, "    _arguments '*: :_files'")
  74         fmt.Fprintln(w, "  ;;")
  75 }
  76
  77 func writeLevel(w io.Writer, root *Command, i int) {
  78         fmt.Fprintf(w, "    case $words[%d] in\n", i)
  79         defer fmt.Fprintln(w, "    esac")
  80
  81         commands := filterByLevel(root, i)
  82         byParent := groupByParent(commands)
  83
  84         for p, c := range byParent {
  85                 names := names(c)
  86                 fmt.Fprintf(w, "      %s)\n", p)
  87                 fmt.Fprintf(w, "        _arguments '%d: :(%s)'\n", i, strings.Join(names, " "))
  88                 fmt.Fprintln(w, "      ;;")
  89         }
  90         fmt.Fprintln(w, "      *)")
  91         fmt.Fprintln(w, "        _arguments '*: :_files'")
  92         fmt.Fprintln(w, "      ;;")
  93
  94 }
  95
  96 func filterByLevel(c *Command, l int) []*Command {
  97         cs := make([]*Command, 0)
  98         if l == 0 {
  99                 cs = append(cs, c)
 100                 return cs
 101         }
 102         for _, s := range c.Commands() {
 103                 cs = append(cs, filterByLevel(s, l-1)...)
 104         }
 105         return cs
 106 }
 107
 108 func groupByParent(commands []*Command) map[string][]*Command {
 109         m := make(map[string][]*Command)
 110         for _, c := range commands {
 111                 parent := c.Parent()
 112                 if parent == nil {
 113                         continue
 114                 }
 115                 m[parent.Name()] = append(m[parent.Name()], c)
 116         }
 117         return m
 118 }
 119
 120 func names(commands []*Command) []string {
 121         ns := make([]string, len(commands))
 122         for i, c := range commands {
 123                 ns[i] = c.Name()
 124         }
 125         return ns
 126 }