Support to install bash completion script automatically (#63)

* Support to install bash completion script automatically

* Bump go from 1.15 to 1.16

.github/workflows/pull-request.yaml

@@ -10,10 +10,10 @@ jobs:
     name: Build
     runs-on: macos-10.15
     steps:
-      - name: Set up Go 1.15
+      - name: Set up Go 1.16
         uses: actions/[email protected]
         with:
-          go-version: 1.15
+          go-version: 1.16
         id: go
       - name: Check out code into the Go module directory
         uses: actions/[email protected]
@@ -29,10 +29,10 @@ jobs:
     name: Lint
     runs-on: ubuntu-latest
     steps:
-      - name: Set up Go 1.15
+      - name: Set up Go 1.16
         uses: actions/[email protected]
         with:
-          go-version: 1.15
+          go-version: 1.16
         id: go
       - name: Check out code into the Go module directory
         uses: actions/[email protected]

.gitignore

@@ -1,4 +1,5 @@
 .idea
 test.xml
+cobra-extension
 
 */***/*.xml

Makefile

@@ -1,6 +1,10 @@
 build: lint fmt
+	go mod tidy
 	go build
 
+copy: build
+	cp cobra-extension /usr/local/bin
+
 lint:
 	golint ./...
 

go.mod

@@ -1,11 +1,12 @@
 module github.com/linuxsuren/cobra-extension
 
-go 1.15
+go 1.16
 
 require (
 	github.com/golang/mock v1.5.0
 	github.com/google/go-github/v29 v29.0.3
 	github.com/linuxsuren/http-downloader v0.0.23
+	github.com/mitchellh/go-homedir v1.1.0
 	github.com/onsi/ginkgo v1.15.1
 	github.com/onsi/gomega v1.11.0
 	github.com/spf13/cobra v1.1.3

main.go

@@ -0,0 +1,19 @@
+package main
+
+import (
+	"github.com/linuxsuren/cobra-extension/pkg"
+	"github.com/spf13/cobra"
+)
+
+func main() {
+	root := &cobra.Command{
+		Use:   "cobra-extension",
+		Short: "This is a demo command for cobra-extension",
+	}
+
+	root.AddCommand(pkg.NewCompletionCmd(root))
+
+	if err := root.Execute(); err != nil {
+		panic(err)
+	}
+}

completion.go → pkg/completion.go

@@ -1,25 +1,29 @@
 package pkg
 
 import (
+	"bytes"
 	"fmt"
+	"github.com/mitchellh/go-homedir"
 	"github.com/spf13/cobra"
+	"os"
+	"path"
 )
 
-// CompletionOptions is the option of completion command
-type CompletionOptions struct {
-	Type string
+// completionOptions is the option of completion command
+type completionOptions struct {
+	shellType string
+	auto      bool
 }
 
 // ShellTypes contains all types of shell
 var ShellTypes = []string{
 	"zsh", "bash", "powerShell",
 }
 
-var completionOptions CompletionOptions
-
 // NewCompletionCmd creates the completion command
 func NewCompletionCmd(rootCmd *cobra.Command) (cmd *cobra.Command) {
 	rootName := rootCmd.Name()
+	opt := completionOptions{}
 
 	cmd = &cobra.Command{
 		Use:   "completion",
@@ -58,25 +62,14 @@ sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/too
 source <(%[1]s completion --type zsh)
 # Set the %[1]s completion code for zsh[1] to autoload on startup
 %[1]s completion --type zsh > "${fpath[1]}/_%[1]s"`, rootName),
-		RunE: func(cmd *cobra.Command, _ []string) (err error) {
-			shellType := completionOptions.Type
-			switch shellType {
-			case "zsh":
-				err = rootCmd.GenZshCompletion(cmd.OutOrStdout())
-			case "powerShell":
-				err = rootCmd.GenPowerShellCompletion(cmd.OutOrStdout())
-			case "bash":
-				err = rootCmd.GenBashCompletion(cmd.OutOrStdout())
-			default:
-				err = fmt.Errorf("unknown shell type %s", shellType)
-			}
-			return
-		},
+		RunE: opt.runE,
 	}
 
 	flags := cmd.Flags()
-	flags.StringVarP(&completionOptions.Type, "type", "", "bash",
+	flags.StringVarP(&opt.shellType, "type", "", "bash",
 		fmt.Sprintf("Generate different types of shell which are %v", ShellTypes))
+	flags.BoolVarP(&opt.auto, "auto", "", true,
+		"Indicate if install the completion script automatically, print the script if it's false")
 
 	err := cmd.RegisterFlagCompletionFunc("type", func(cmd *cobra.Command, args []string, toComplete string) (
 		i []string, directive cobra.ShellCompDirective) {
@@ -87,3 +80,72 @@ source <(%[1]s completion --type zsh)
 	}
 	return
 }
+
+func (o *completionOptions) runE(cmd *cobra.Command, _ []string) (err error) {
+	rootCmd := cmd.Root()
+	writer := cmd.OutOrStdout()
+	buf := bytes.NewBuffer([]byte{})
+	if o.auto {
+		writer = buf
+	}
+
+	shellType := o.shellType
+	switch shellType {
+	case "zsh":
+		err = rootCmd.GenZshCompletion(writer)
+	case "powerShell":
+		err = rootCmd.GenPowerShellCompletion(writer)
+	case "bash":
+		err = rootCmd.GenBashCompletion(writer)
+	default:
+		err = fmt.Errorf("unknown shell type %s", shellType)
+	}
+
+	if err == nil && o.auto {
+		err = installCompletionScript(rootCmd.Use, buf.Bytes())
+	}
+	return
+}
+
+func installCompletionScript(name string, script []byte) (err error) {
+	// create ~/.bash_completion if it does not exist or is empty
+	if err = writeIfNotExist("~/.bash_completion", []byte(`
+for bcfile in ~/.bash_completion.d/* ; do
+  . $bcfile
+done`), false); err != nil {
+		return
+	}
+
+	// install script for the current command to '~/.bash_completion.d/cmd'
+	if err = writeIfNotExist(fmt.Sprintf("~/.bash_completion.d/%s", name), script, true); err != nil {
+		return
+	}
+	return
+}
+
+func writeIfNotExist(requestPath string, data []byte, force bool) (err error) {
+	var targetPath string
+	if targetPath, err = homedir.Expand(requestPath); err != nil {
+		return
+	}
+
+	var ok bool
+	if ok, err = pathExists(targetPath); !ok || err != nil || force {
+		if err = os.MkdirAll(path.Dir(targetPath), os.FileMode(0755)); err != nil {
+			return
+		}
+		err = os.WriteFile(targetPath, data, 0644)
+	}
+	return
+}
+
+func pathExists(path string) (bool, error) {
+	_, err := os.Stat(path)
+	if err == nil {
+		return true, nil
+	}
+	if os.IsNotExist(err) {
+		return false, nil
+	}
+	return false, err
+}

padding_test.go → pkg/padding_test.go

@@ -2,7 +2,6 @@ package pkg
 
 import (
 	"github.com/golang/mock/gomock"
-	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 )