V2

1 files changed, 296 insertions, 68 deletions

diff --git a/complete.go b/complete.go
index 423cbec..7e4a3d5 100644
--- a/complete.go
+++ b/complete.go
@@ -1,104 +1,332 @@
 package complete
 
 import (
-	"flag"
 	"fmt"
 	"io"
 	"os"
 	"strconv"
 	"strings"
 
-	"github.com/posener/complete/cmd"
+	"github.com/posener/complete/internal/arg"
+	"github.com/posener/complete/internal/install"
+	"github.com/posener/complete/internal/tokener"
 )
 
-const (
-	envLine  = "COMP_LINE"
-	envPoint = "COMP_POINT"
-	envDebug = "COMP_DEBUG"
+// Completer is an interface that a command line should implement in order to get bash completion.
+type Completer interface {
+	// SubCmdList should return the list of all sub commands of the current command.
+	SubCmdList() []string
+	// SubCmdGet should return a sub command of the current command for the given sub command name.
+	SubCmdGet(cmd string) Completer
+	// FlagList should return a list of all the flag names of the current command. The flag names
+	// should not have the dash prefix.
+	FlagList() []string
+	// FlagGet should return completion options for a given flag. It is invoked with the flag name
+	// without the dash prefix. The flag is not promised to be in the command flags. In that case,
+	// this method should return a nil predictor.
+	FlagGet(flag string) Predictor
+	// ArgsGet should return predictor for positional arguments of the command line.
+	ArgsGet() Predictor
+}
+
+// Predictor can predict completion options.
+type Predictor interface {
+	// Predict returns prediction options for a given prefix. The prefix is what currently is typed
+	// as a hint for what to return, but the returned values can have any prefix. The returned
+	// values will be filtered by the prefix when needed regardless. The prefix may be empty which
+	// means that no value was typed.
+	Predict(prefix string) []string
+}
+
+// PredictFunc is a function that implements the Predictor interface.
+type PredictFunc func(prefix string) []string
+
+func (p PredictFunc) Predict(prefix string) []string {
+	if p == nil {
+		return nil
+	}
+	return p(prefix)
+}
+
+var (
+	getEnv           = os.Getenv
+	exit             = os.Exit
+	out    io.Writer = os.Stdout
+	in     io.Reader = os.Stdin
 )
 
-// Complete structs define completion for a command with CLI options
-type Complete struct {
-	Command Command
-	cmd.CLI
-	Out io.Writer
+// Complete the command line arguments for the given command in the case that the program
+// was invoked with COMP_LINE and COMP_POINT environment variables. In that case it will also
+// `os.Exit()`. The program name should be provided for installation purposes.
+func Complete(name string, cmd Completer) {
+	var (
+		line        = getEnv("COMP_LINE")
+		point       = getEnv("COMP_POINT")
+		doInstall   = getEnv("COMP_INSTALL") == "1"
+		doUninstall = getEnv("COMP_UNINSTALL") == "1"
+		yes         = getEnv("COMP_YES") == "1"
+	)
+	if doInstall || doUninstall {
+		install.Run(name, doUninstall, yes, out, in)
+		exit(0)
+		return
+	}
+	if line == "" {
+		return
+	}
+	i, err := strconv.Atoi(point)
+	if err != nil {
+		panic("COMP_POINT env should be integer, got: " + point)
+	}
+
+	// Parse the command line up to the completion point.
+	args := arg.Parse(line[:i])
+
+	// The first word is the current command name.
+	args = args[1:]
+
+	// Run the completion algorithm.
+	options, err := completer{Completer: cmd, args: args}.complete()
+	if err != nil {
+		fmt.Fprintln(out, "\n"+err.Error())
+	} else {
+		for _, option := range options {
+			fmt.Fprintln(out, option)
+		}
+	}
+	exit(0)
+}
+
+type completer struct {
+	Completer
+	args  []arg.Arg
+	stack []Completer
+}
+
+// compete command with given before and after text.
+// if the command has sub commands: try to complete only sub commands or help flags. Otherwise
+// complete flags and positional arguments.
+func (c completer) complete() ([]string, error) {
+reset:
+	arg := arg.Arg{}
+	if len(c.args) > 0 {
+		arg = c.args[0]
+	}
+	switch {
+	case len(c.SubCmdList()) == 0:
+		// No sub commands, parse flags and positional arguments.
+		return c.suggestLeafCommandOptions(), nil
+
+	// case !arg.Completed && arg.IsFlag():
+	// Suggest help flags for command
+	// return []string{helpFlag(arg.Text)}, nil
+
+	case !arg.Completed:
+		// Currently typing a sub command.
+		return c.suggestSubCommands(arg.Text), nil
+
+	case c.SubCmdGet(arg.Text) != nil:
+		// Sub command completed, look into that sub command completion.
+		// Set the complete command to the requested sub command, and the before text to all the text
+		// after the command name and rerun the complete algorithm with the new sub command.
+		c.stack = append([]Completer{c.Completer}, c.stack...)
+		c.Completer = c.SubCmdGet(arg.Text)
+		c.args = c.args[1:]
+		goto reset
+
+	default:
+
+		// Sub command is unknown...
+		return nil, fmt.Errorf("unknown subcommand: %s", arg.Text)
+	}
+}
+
+func (c completer) suggestSubCommands(prefix string) []string {
+	if len(prefix) > 0 && prefix[0] == '-' {
+		return []string{helpFlag(prefix)}
+	}
+	subs := c.SubCmdList()
+	return suggest("", prefix, func(prefix string) []string {
+		var options []string
+		for _, sub := range subs {
+			if strings.HasPrefix(sub, prefix) {
+				options = append(options, sub)
+			}
+		}
+		return options
+	})
+}
+
+func (c completer) suggestLeafCommandOptions() (options []string) {
+	arg, before := arg.Arg{}, arg.Arg{}
+	if len(c.args) > 0 {
+		arg = c.args[len(c.args)-1]
+	}
+	if len(c.args) > 1 {
+		before = c.args[len(c.args)-2]
+	}
+
+	if !arg.Completed {
+		// Complete value being typed.
+		if arg.HasValue {
+			// Complete value of current flag.
+			if arg.HasFlag {
+				return c.suggestFlagValue(arg.Flag, arg.Value)
+			}
+			// Complete value of flag in a previous argument.
+			if before.HasFlag && !before.HasValue {
+				return c.suggestFlagValue(before.Flag, arg.Value)
+			}
+		}
+
+		// A value with no flag. Suggest positional argument.
+		if !arg.HasValue {
+			options = c.suggestFlag(arg.Dashes, arg.Flag)
+		}
+		if !arg.HasFlag {
+			options = append(options, c.suggestArgsValue(arg.Value)...)
+		}
+		// Suggest flag according to prefix.
+		return options
+	}
+
+	// Has a value that was already completed. Suggest all flags and positional arguments.
+	if arg.HasValue {
+		options = c.suggestFlag(arg.Dashes, "")
+		if !arg.HasFlag {
+			options = append(options, c.suggestArgsValue("")...)
+		}
+		return options
+	}
+	// A flag without a value. Suggest a value or suggest any flag.
+	options = c.suggestFlagValue(arg.Flag, "")
+	if len(options) > 0 {
+		return options
+	}
+	return c.suggestFlag("", "")
 }
 
-// New creates a new complete command.
-// name is the name of command we want to auto complete.
-// IMPORTANT: it must be the same name - if the auto complete
-// completes the 'go' command, name must be equal to "go".
-// command is the struct of the command completion.
-func New(name string, command Command) *Complete {
-	return &Complete{
-		Command: command,
-		CLI:     cmd.CLI{Name: name},
-		Out:     os.Stdout,
+func (c completer) suggestFlag(dashes, prefix string) []string {
+	if dashes == "" {
+		dashes = "-"
 	}
+	return suggest(dashes, prefix, func(prefix string) []string {
+		var options []string
+		c.iterateStack(func(cmd Completer) {
+			// Suggest all flags with the given prefix.
+			for _, name := range cmd.FlagList() {
+				if strings.HasPrefix(name, prefix) {
+					options = append(options, dashes+name)
+				}
+			}
+		})
+		return options
+	})
+}
+
+func (c completer) suggestFlagValue(flagName, prefix string) []string {
+	var options []string
+	c.iterateStack(func(cmd Completer) {
+		if len(options) == 0 {
+			if p := cmd.FlagGet(flagName); p != nil {
+				options = p.Predict(prefix)
+			}
+		}
+	})
+	return filterByPrefix(prefix, options...)
+}
+
+func (c completer) suggestArgsValue(prefix string) []string {
+	var options []string
+	c.iterateStack(func(cmd Completer) {
+		if len(options) == 0 {
+			if p := cmd.ArgsGet(); p != nil {
+				options = p.Predict(prefix)
+			}
+		}
+	})
+	return filterByPrefix(prefix, options...)
 }
 
-// Run runs the completion and add installation flags beforehand.
-// The flags are added to the main flag CommandLine variable.
-func (c *Complete) Run() bool {
-	c.AddFlags(nil)
-	flag.Parse()
-	return c.Complete()
+func (c completer) iterateStack(f func(Completer)) {
+	for _, cmd := range append([]Completer{c.Completer}, c.stack...) {
+		f(cmd)
+	}
 }
 
-// Complete a command from completion line in environment variable,
-// and print out the complete options.
-// returns success if the completion ran or if the cli matched
-// any of the given flags, false otherwise
-// For installation: it assumes that flags were added and parsed before
-// it was called.
-func (c *Complete) Complete() bool {
-	line, point, ok := getEnv()
-	if !ok {
-		// make sure flags parsed,
-		// in case they were not added in the main program
-		return c.CLI.Run()
+func suggest(dashes, prefix string, collect func(prefix string) []string) []string {
+	options := collect(prefix)
+	// If nothing was suggested, suggest all flags.
+	if len(options) == 0 {
+		prefix = ""
+		options = collect(prefix)
 	}
 
-	if point >= 0 && point < len(line) {
-		line = line[:point]
+	// Add help flag if needed.
+	help := helpFlag(dashes + prefix)
+	if len(options) == 0 || strings.HasPrefix(help, dashes+prefix) {
+		options = append(options, help)
 	}
 
-	Log("Completing phrase: %s", line)
-	a := newArgs(line)
-	Log("Completing last field: %s", a.Last)
-	options := c.Command.Predict(a)
-	Log("Options: %s", options)
+	return options
+}
 
-	// filter only options that match the last argument
-	matches := []string{}
+func filterByPrefix(prefix string, options ...string) []string {
+	var filtered []string
 	for _, option := range options {
-		if strings.HasPrefix(option, a.Last) {
-			matches = append(matches, option)
+		if fixed, ok := hasPrefix(option, prefix); ok {
+			filtered = append(filtered, fixed)
 		}
 	}
-	Log("Matches: %s", matches)
-	c.output(matches)
-	return true
+	if len(filtered) > 0 {
+		return filtered
+	}
+	return options
 }
 
-func getEnv() (line string, point int, ok bool) {
-	line = os.Getenv(envLine)
-	if line == "" {
-		return
+// hasPrefix checks if s has the give prefix. It disregards quotes and escaped spaces, and return
+// s in the form of the given prefix.
+func hasPrefix(s, prefix string) (string, bool) {
+	var (
+		token  tokener.Tokener
+		si, pi int
+	)
+	for ; pi < len(prefix); pi++ {
+		token.Visit(prefix[pi])
+		lastQuote := !token.Escaped() && (prefix[pi] == '"' || prefix[pi] == '\'')
+		if lastQuote {
+			continue
+		}
+		if si == len(s) {
+			break
+		}
+		if s[si] == ' ' && !token.Quoted() && token.Escaped() {
+			s = s[:si] + "\\" + s[si:]
+		}
+		if s[si] != prefix[pi] {
+			return "", false
+		}
+		si++
 	}
-	point, err := strconv.Atoi(os.Getenv(envPoint))
-	if err != nil {
-		// If failed parsing point for some reason, set it to point
-		// on the end of the line.
-		Log("Failed parsing point %s: %v", os.Getenv(envPoint), err)
-		point = len(line)
+
+	if pi < len(prefix) {
+		return "", false
 	}
-	return line, point, true
+
+	for ; si < len(s); si++ {
+		token.Visit(s[si])
+	}
+
+	return token.Closed(), true
 }
 
-func (c *Complete) output(options []string) {
-	// stdout of program defines the complete options
-	for _, option := range options {
-		fmt.Fprintln(c.Out, option)
+// helpFlag returns either "-h", "-help" or "--help".
+func helpFlag(prefix string) string {
+	if prefix == "" || prefix == "-" || prefix == "-h" {
+		return "-h"
+	}
+	if strings.HasPrefix(prefix, "--") {
+		return "--help"
 	}
+	return "-help"
 }