Merge pull request #18 from posener/zsh

zsh support

cmd/install/bash.go

@@ -1,153 +1,32 @@
 package install
 
-import (
-	"bufio"
-	"errors"
-	"fmt"
-	"io"
-	"io/ioutil"
-	"os"
-	"os/user"
-	"path/filepath"
-)
+import "fmt"
 
-type bash struct{}
-
-func (bash) Install(cmd, bin string) error {
-	bashRCFileName, err := bashRCFileName()
-	if err != nil {
-		return err
-	}
-	completeCmd := completeCmd(cmd, bin)
-	if isInFile(bashRCFileName, completeCmd) {
-		return errors.New("Already installed in ~/.bashrc")
-	}
-
-	bashRC, err := os.OpenFile(bashRCFileName, os.O_RDWR|os.O_APPEND, 0)
-	if err != nil {
-		return err
-	}
-	defer bashRC.Close()
-	_, err = bashRC.WriteString(fmt.Sprintf("\n%s\n", completeCmd))
-	return err
+// (un)install in bash
+// basically adds/remove from .bashrc:
+//
+// complete -C </path/to/completion/command> <command>
+type bash struct {
+	rc string
 }
 
-func (bash) Uninstall(cmd, bin string) error {
-	bashRC, err := bashRCFileName()
-	if err != nil {
-		return err
+func (b bash) Install(cmd, bin string) error {
+	completeCmd := b.cmd(cmd, bin)
+	if lineInFile(b.rc, completeCmd) {
+		return fmt.Errorf("already installed in %s", b.rc)
 	}
-	backup := bashRC + ".bck"
-	err = copyFile(bashRC, backup)
-	if err != nil {
-		return err
-	}
-	completeCmd := completeCmd(cmd, bin)
-	if !isInFile(bashRC, completeCmd) {
-		return errors.New("Does not installed in ~/.bashrc")
-	}
-	temp, err := uninstallToTemp(bashRC, completeCmd)
-	if err != nil {
-		return err
-	}
-
-	err = copyFile(temp, bashRC)
-	if err != nil {
-		return err
-	}
-
-	return os.Remove(backup)
-
+	return appendToFile(b.rc, completeCmd)
 }
 
-func completeCmd(cmd, bin string) string {
+func (b bash) Uninstall(cmd, bin string) error {
+	completeCmd := b.cmd(cmd, bin)
+	if !lineInFile(b.rc, completeCmd) {
+		return fmt.Errorf("does not installed in %s", b.rc)
+	}
+
+	return removeFromFile(b.rc, completeCmd)
+}
+
+func (bash) cmd(cmd, bin string) string {
 	return fmt.Sprintf("complete -C %s %s", bin, cmd)
 }
-
-func bashRCFileName() (string, error) {
-	u, err := user.Current()
-	if err != nil {
-		return "", err
-	}
-	return filepath.Join(u.HomeDir, ".bashrc"), nil
-}
-
-func isInFile(name string, lookFor string) bool {
-	f, err := os.Open(name)
-	if err != nil {
-		return false
-	}
-	defer f.Close()
-	r := bufio.NewReader(f)
-	prefix := []byte{}
-	for {
-		line, isPrefix, err := r.ReadLine()
-		if err == io.EOF {
-			return false
-		}
-		if err != nil {
-			return false
-		}
-		if isPrefix {
-			prefix = append(prefix, line...)
-			continue
-		}
-		line = append(prefix, line...)
-		if string(line) == lookFor {
-			return true
-		}
-		prefix = prefix[:0]
-	}
-}
-
-func uninstallToTemp(bashRCFileName, completeCmd string) (string, error) {
-	rf, err := os.Open(bashRCFileName)
-	if err != nil {
-		return "", err
-	}
-	defer rf.Close()
-	wf, err := ioutil.TempFile("/tmp", "bashrc-")
-	if err != nil {
-		return "", err
-	}
-	defer wf.Close()
-
-	r := bufio.NewReader(rf)
-	prefix := []byte{}
-	for {
-		line, isPrefix, err := r.ReadLine()
-		if err == io.EOF {
-			break
-		}
-		if err != nil {
-			return "", err
-		}
-		if isPrefix {
-			prefix = append(prefix, line...)
-			continue
-		}
-		line = append(prefix, line...)
-		str := string(line)
-		if str == completeCmd {
-			continue
-		}
-		wf.WriteString(str + "\n")
-		prefix = prefix[:0]
-	}
-	return wf.Name(), nil
-}
-
-func copyFile(src string, dst string) error {
-	in, err := os.Open(src)
-	if err != nil {
-		return err
-	}
-	defer in.Close()
-	out, err := os.Create(dst)
-	if err != nil {
-		return err
-	}
-	defer out.Close()
-	_, err = io.Copy(out, in)
-	return err
-}

cmd/install/install.go

@@ -2,9 +2,11 @@ package install
 
 import (
 	"errors"
-	"fmt"
 	"os"
+	"os/user"
 	"path/filepath"
+
+	"github.com/hashicorp/go-multierror"
 )
 
 type installer interface {
@@ -15,46 +17,55 @@ type installer interface {
 // Install complete command given:
 // cmd: is the command name
 func Install(cmd string) error {
-	shell := shellType()
-	if shell == "" {
-		return errors.New("must install through a terminatl")
-	}
-	i := getInstaller(shell)
-	if i == nil {
-		return fmt.Errorf("shell %s not supported", shell)
+	is := installers()
+	if len(is) == 0 {
+		return errors.New("Did not found any shells to install")
 	}
 	bin, err := getBinaryPath()
 	if err != nil {
 		return err
 	}
-	return i.Install(cmd, bin)
+
+	for _, i := range is {
+		errI := i.Install(cmd, bin)
+		if errI != nil {
+			multierror.Append(err, errI)
+		}
+	}
+
+	return err
 }
 
 // Uninstall complete command given:
 // cmd: is the command name
 func Uninstall(cmd string) error {
-	shell := shellType()
-	if shell == "" {
-		return errors.New("must uninstall through a terminatl")
-	}
-	i := getInstaller(shell)
-	if i == nil {
-		return fmt.Errorf("shell %s not supported", shell)
+	is := installers()
+	if len(is) == 0 {
+		return errors.New("Did not found any shells to uninstall")
 	}
 	bin, err := getBinaryPath()
 	if err != nil {
 		return err
 	}
-	return i.Uninstall(cmd, bin)
+
+	for _, i := range is {
+		errI := i.Uninstall(cmd, bin)
+		if errI != nil {
+			multierror.Append(err, errI)
+		}
+	}
+
+	return err
 }
 
-func getInstaller(shell string) installer {
-	switch shell {
-	case "bash":
-		return bash{}
-	default:
-		return nil
+func installers() (i []installer) {
+	if f := rcFile(".bashrc"); f != "" {
+		i = append(i, bash{f})
 	}
+	if f := rcFile(".zshrc"); f != "" {
+		i = append(i, zsh{f})
+	}
+	return
 }
 
 func getBinaryPath() (string, error) {
@@ -65,7 +76,10 @@ func getBinaryPath() (string, error) {
 	return filepath.Abs(bin)
 }
 
-func shellType() string {
-	shell := os.Getenv("SHELL")
-	return filepath.Base(shell)
+func rcFile(name string) string {
+	u, err := user.Current()
+	if err != nil {
+		return ""
+	}
+	return filepath.Join(u.HomeDir, name)
 }

cmd/install/utils.go

@@ -0,0 +1,118 @@
+package install
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+)
+
+func lineInFile(name string, lookFor string) bool {
+	f, err := os.Open(name)
+	if err != nil {
+		return false
+	}
+	defer f.Close()
+	r := bufio.NewReader(f)
+	prefix := []byte{}
+	for {
+		line, isPrefix, err := r.ReadLine()
+		if err == io.EOF {
+			return false
+		}
+		if err != nil {
+			return false
+		}
+		if isPrefix {
+			prefix = append(prefix, line...)
+			continue
+		}
+		line = append(prefix, line...)
+		if string(line) == lookFor {
+			return true
+		}
+		prefix = prefix[:0]
+	}
+}
+
+func appendToFile(name string, content string) error {
+	f, err := os.OpenFile(name, os.O_RDWR|os.O_APPEND, 0)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	_, err = f.WriteString(fmt.Sprintf("\n%s\n", content))
+	return err
+}
+
+func removeFromFile(name string, content string) error {
+	backup := name + ".bck"
+	err := copyFile(name, backup)
+	if err != nil {
+		return err
+	}
+	temp, err := removeContentToTempFile(name, content)
+	if err != nil {
+		return err
+	}
+
+	err = copyFile(temp, name)
+	if err != nil {
+		return err
+	}
+
+	return os.Remove(backup)
+}
+
+func removeContentToTempFile(name, content string) (string, error) {
+	rf, err := os.Open(name)
+	if err != nil {
+		return "", err
+	}
+	defer rf.Close()
+	wf, err := ioutil.TempFile("/tmp", "complete-")
+	if err != nil {
+		return "", err
+	}
+	defer wf.Close()
+
+	r := bufio.NewReader(rf)
+	prefix := []byte{}
+	for {
+		line, isPrefix, err := r.ReadLine()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return "", err
+		}
+		if isPrefix {
+			prefix = append(prefix, line...)
+			continue
+		}
+		line = append(prefix, line...)
+		str := string(line)
+		if str == content {
+			continue
+		}
+		wf.WriteString(str + "\n")
+		prefix = prefix[:0]
+	}
+	return wf.Name(), nil
+}
+
+func copyFile(src string, dst string) error {
+	in, err := os.Open(src)
+	if err != nil {
+		return err
+	}
+	defer in.Close()
+	out, err := os.Create(dst)
+	if err != nil {
+		return err
+	}
+	defer out.Close()
+	_, err = io.Copy(out, in)
+	return err
+}

cmd/install/zsh.go

@@ -0,0 +1,39 @@
+package install
+
+import "fmt"
+
+// (un)install in zsh
+// basically adds/remove from .zshrc:
+//
+// autoload -U +X bashcompinit && bashcompinit"
+// complete -C </path/to/completion/command> <command>
+type zsh struct {
+	rc string
+}
+
+func (z zsh) Install(cmd, bin string) error {
+	completeCmd := z.cmd(cmd, bin)
+	if lineInFile(z.rc, completeCmd) {
+		return fmt.Errorf("already installed in %s", z.rc)
+	}
+
+	bashCompInit := "autoload -U +X bashcompinit && bashcompinit"
+	if !lineInFile(z.rc, bashCompInit) {
+		completeCmd = bashCompInit + "\n" + completeCmd
+	}
+
+	return appendToFile(z.rc, completeCmd)
+}
+
+func (z zsh) Uninstall(cmd, bin string) error {
+	completeCmd := z.cmd(cmd, bin)
+	if !lineInFile(z.rc, completeCmd) {
+		return fmt.Errorf("does not installed in %s", z.rc)
+	}
+
+	return removeFromFile(z.rc, completeCmd)
+}
+
+func (zsh) cmd(cmd, bin string) string {
+	return fmt.Sprintf("complete -C %s %s", bin, cmd)
+}

gocomplete/complete.go

@@ -4,11 +4,11 @@ package main
 import "github.com/posener/complete"
 
 var (
-	ellipsis     = complete.PredictSet("./...")
-	anyPackage   = predictPackages("")
-	goFiles      = complete.PredictFiles("*.go")
-	anyFile      = complete.PredictFiles("*")
-	anyGo        = complete.PredictOr(goFiles, anyPackage, ellipsis)
+	ellipsis   = complete.PredictSet("./...")
+	anyPackage = predictPackages("")
+	goFiles    = complete.PredictFiles("*.go")
+	anyFile    = complete.PredictFiles("*")
+	anyGo      = complete.PredictOr(goFiles, anyPackage, ellipsis)
 )
 
 func main() {

readme.md

@@ -36,7 +36,14 @@ Uninstall by `gocomplete -uninstall`
 - Complete packages names or `.go` files when necessary.
 - Complete test names after `-run` flag.
 
-## Usage
+## complete package
+
+Supported shells:
+
+[x] bash
+[x] zsh
+
+### Usage
 
 Assuming you have program called `run` and you want to have bash completion
 for it, meaning, if you type `run` then space, then press the `Tab` key,
@@ -95,7 +102,7 @@ func main() {
 }
 ```
 
-## Self completing program
+### Self completing program
 
 In case that the program that we want to complete is written in go we
 can make it self completing. Here is an [example](./example/self/main.go)