feature: Add `goctl completion` (#1505)

* feature: Add `goctl completion`

* Update const

Co-authored-by: anqiansong <anqiansong@bytedance.com>

tools/goctl/completion/completion.go

@@ -0,0 +1,75 @@
+package completion
+
+import (
+	"bytes"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"runtime"
+
+	"github.com/logrusorgru/aurora"
+	"github.com/urfave/cli"
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
+	"github.com/zeromicro/go-zero/tools/goctl/vars"
+)
+
+func Completion(c *cli.Context) error {
+	goos := runtime.GOOS
+	if goos == vars.OsWindows {
+		return fmt.Errorf("%q: only support unix-like OS", goos)
+	}
+	name := c.String("name")
+	if len(name) == 0 {
+		name = defaultCompletionFilename
+	}
+	if filepath.IsAbs(name) {
+		return fmt.Errorf("unsupport absolute path: %q", name)
+	}
+
+	home, err := pathx.GetAutoCompleteHome()
+	if err != nil {
+		return err
+	}
+	buffer := bytes.NewBuffer(nil)
+	zshF := filepath.Join(home, "zsh", defaultCompletionFilename)
+	err = pathx.MkdirIfNotExist(filepath.Dir(zshF))
+	if err != nil {
+		return err
+	}
+
+	bashF := filepath.Join(home, "bash", defaultCompletionFilename)
+	err = pathx.MkdirIfNotExist(filepath.Dir(bashF))
+	if err != nil {
+		return err
+	}
+	var flag = magic
+	err = ioutil.WriteFile(zshF, zsh, os.ModePerm)
+	if err != nil {
+		return err
+	}
+
+	flag |= flagZsh
+	err = ioutil.WriteFile(bashF, bash, os.ModePerm)
+	if err != nil {
+		return err
+	}
+
+	flag |= flagBash
+	buffer.WriteString(aurora.Green("generation auto completion success!\n").String())
+	buffer.WriteString(aurora.Green("executes the following script to setting shell:\n").String())
+	switch flag {
+	case magic | flagZsh:
+		buffer.WriteString(aurora.Blue(fmt.Sprintf("echo PROG=goctl source %s >> ~/.zshrc && source ~/.zshrc", zshF)).String())
+	case magic | flagBash:
+		buffer.WriteString(aurora.Blue(fmt.Sprintf("echo PROG=goctl source %s >> ~/.bashrc && source ~/.bashrc", bashF)).String())
+	case magic | flagZsh | flagBash:
+		buffer.WriteString(aurora.Blue(fmt.Sprintf(`echo PROG=goctl source %s >> ~/.zshrc && source ~/.zshrc
+or
+echo PROG=goctl source %s >> ~/.bashrc && source ~/.bashrc`, zshF, bashF)).String())
+	default:
+		return nil
+	}
+	fmt.Println(buffer.String())
+	return nil
+}

tools/goctl/completion/const.go

@@ -0,0 +1,9 @@
+package completion
+
+const BashCompletionFlag = `generate-goctl-completion`
+const defaultCompletionFilename = "goctl_autocomplete"
+const (
+	magic = 1 << iota
+	flagZsh
+	flagBash
+)

tools/goctl/completion/script.go

@@ -0,0 +1,51 @@
+package completion
+
+import "fmt"
+
+var zsh = []byte(fmt.Sprintf(`#compdef $PROG
+
+_cli_zsh_autocomplete() {
+
+  local -a opts
+  local cur
+  cur=${words[-1]}
+  if [[ "$cur" == "-"* ]]; then
+    opts=("${(@f)$(_CLI_ZSH_AUTOCOMPLETE_HACK=1 ${words[@]:0:#words[@]-1} ${cur} --%s)}")
+  else
+    opts=("${(@f)$(_CLI_ZSH_AUTOCOMPLETE_HACK=1 ${words[@]:0:#words[@]-1} --%s)}")
+  fi
+
+  if [[ "${opts[1]}" != "" ]]; then
+    _describe 'values' opts
+  else
+    _files
+  fi
+
+  return
+}
+
+compdef _cli_zsh_autocomplete $PROG
+`, BashCompletionFlag, BashCompletionFlag))
+
+var bash = []byte(fmt.Sprintf(`#! /bin/bash
+
+: ${PROG:=$(basename ${BASH_SOURCE})}
+
+_cli_bash_autocomplete() {
+  if [[ "${COMP_WORDS[0]}" != "source" ]]; then
+    local cur opts base
+    COMPREPLY=()
+    cur="${COMP_WORDS[COMP_CWORD]}"
+    if [[ "$cur" == "-"* ]]; then
+      opts=$( ${COMP_WORDS[@]:0:$COMP_CWORD} ${cur} --%s )
+    else
+      opts=$( ${COMP_WORDS[@]:0:$COMP_CWORD} --%s )
+    fi
+    COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
+    return 0
+  fi
+}
+
+complete -o bashdefault -o default -o nospace -F _cli_bash_autocomplete $PROG
+unset PROG
+`, BashCompletionFlag, BashCompletionFlag))

tools/goctl/goctl.go

@@ -20,6 +20,7 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/api/tsgen"
 	"github.com/zeromicro/go-zero/tools/goctl/api/validate"
 	"github.com/zeromicro/go-zero/tools/goctl/bug"
+	"github.com/zeromicro/go-zero/tools/goctl/completion"
 	"github.com/zeromicro/go-zero/tools/goctl/docker"
 	"github.com/zeromicro/go-zero/tools/goctl/internal/errorx"
 	"github.com/zeromicro/go-zero/tools/goctl/internal/version"
@@ -797,13 +798,29 @@ var commands = []cli.Command{
 			},
 		},
 	},
+	{
+		Name:   "completion",
+		Usage:  "generation completion script, it only works for unix-like OS",
+		Action: completion.Completion,
+		Flags: []cli.Flag{
+			cli.StringFlag{
+				Name:  "name, n",
+				Usage: "the filename of auto complete script, default is [goctl_autocomplete]",
+			},
+		},
+	},
 }
 
 func main() {
 	logx.Disable()
 	load.Disable()
 
+	cli.BashCompletionFlag = cli.BoolFlag{
+		Name:   completion.BashCompletionFlag,
+		Hidden: true,
+	}
 	app := cli.NewApp()
+	app.EnableBashCompletion = true
 	app.Usage = "a cli tool to generate code"
 	app.Version = fmt.Sprintf("%s %s/%s", version.BuildVersion, runtime.GOOS, runtime.GOARCH)
 	app.Commands = commands

tools/goctl/util/pathx/file.go

@@ -18,6 +18,7 @@ const (
 	NL              = "\n"
 	goctlDir        = ".goctl"
 	gitDir          = ".git"
+	autoCompleteDir = ".auto_complete"
 )
 
 var goctlHome string
@@ -93,6 +94,16 @@ func GetGitHome() (string, error) {
 	return filepath.Join(goctlH, gitDir), nil
 }
 
+// GetAutoCompleteHome returns the auto_complete home of goctl.
+func GetAutoCompleteHome() (string, error) {
+	goctlH, err := GetGoctlHome()
+	if err != nil {
+		return "", err
+	}
+
+	return filepath.Join(goctlH, autoCompleteDir), nil
+}
+
 // GetTemplateDir returns the category path value in GoctlHome where could get it by GetGoctlHome
 func GetTemplateDir(category string) (string, error) {
 	home, err := GetGoctlHome()