micro/gateway/cli/executor.go

package cli

import (
	"errors"
	"fmt"
	"git.nspix.com/golang/micro/helper/console"
	"git.nspix.com/golang/micro/helper/unsafestr"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

var (
	_seq int64
)

type Executor struct {
	name           string
	parent         *Executor
	usage          string
	description    string
	handleFunc     HandleFunc
	children       map[string]*Executor
	createdTime    time.Time
	seq            int64
	locker         sync.RWMutex
	params         []string //params (eg: /show/user/:uid => :uid)
	NotFoundHandle func(ctx *Context, cmd string) (*Response, error)
}

type (
	executors []*Executor
)

func (p executors) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p executors) Len() int           { return len(p) }
func (p executors) Less(i, j int) bool { return p[i].seq < p[j].seq }

func (exec *Executor) hasChildren(name string) bool {
	exec.locker.RLock()
	defer exec.locker.RUnlock()
	if name == "" {
		return len(exec.children) > 0
	}
	if _, ok := exec.children[name]; ok {
		return true
	}
	return false
}

func (exec *Executor) getDescription() string {
	if exec.description != "" {
		return exec.description
	}
	var p *Executor
	var text string
	p = exec
	for {
		if p.parent == nil {
			break
		}
		text = p.name + " " + text
		p = p.parent
	}
	return text
}

func (exec *Executor) Children(name string) (*Executor, error) {
	exec.locker.RLock()
	defer exec.locker.RUnlock()
	if cmd, ok := exec.children[name]; ok {
		return cmd, nil
	}
	return nil, errors.New("children not found")
}

func (exec *Executor) Completer(tokens ...string) []string {
	exec.locker.RLock()
	defer exec.locker.RUnlock()
	p := exec.children
	length := len(tokens)
	var token, prefix string
	result := make([]string, 0)
	if length > 0 {
		for i := 0; i < length; i++ {
			token = tokens[i]
			if i == length-1 {
				for k, e := range p {
					if token == k {
						if e.children != nil {
							for kk, _ := range e.children {
								result = append(result, prefix+token+" "+kk)
							}
						}
					} else if strings.HasPrefix(k, token) {
						result = append(result, prefix+k)
					}
				}
			} else {
				if next, ok := p[token]; !ok || next.children == nil {
					return []string{}
				} else {
					p = next.children
				}
			}
			prefix += token + " "
		}
	} else {
		for k, _ := range p {
			result = append(result, k)
		}
	}
	return result
}

func (exec *Executor) String() string {
	table := console.NewTable()
	var loop func(string, *Executor)
	loop = func(prefix string, e *Executor) {
		if e.children == nil {
			return
		}
		vs := make(executors, len(e.children))
		var i int
		for _, v := range e.children {
			vs[i] = v
			i++
		}
		sort.Sort(vs)
		for _, v := range vs {
			if prefix == "" {
				if v.handleFunc != nil {
					table.AddRow(v.name, v.getDescription())
				}
			} else {
				if v.handleFunc != nil {
					table.AddRow(prefix+" "+v.name, v.getDescription())
				}
			}
			if prefix == "" {
				loop(v.name, v)
			} else {
				loop(prefix+" "+v.name, v)
			}
		}
	}
	loop("", exec)
	buf, _ := table.Marshal()
	return unsafestr.BytesToString(buf)
}

func (exec *Executor) Append(child ...*Executor) *Executor {
	exec.locker.Lock()
	defer exec.locker.Unlock()
	for _, v := range child {
		v.parent = exec
		exec.children[v.name] = v
	}
	return exec
}

func (exec *Executor) WithHandle(cb HandleFunc) *Executor {
	exec.handleFunc = cb
	return exec
}

func (exec *Executor) Do(ctx *Context, args ...string) (res *Response, err error) {
	var (
		root interface{}
	)
	if root = ctx.Get("ROOT"); root == nil {
		ctx.Set("ROOT", exec)
	}
	ctx.Set("arguments", args)
	if len(args) > 0 && exec.hasChildren(args[0]) {
		if len(args) > 1 {
			return exec.children[args[0]].Do(ctx, args[1:]...)
		} else {
			return exec.children[args[0]].Do(ctx)
		}
	}
	if exec.handleFunc == nil {
		if exec.NotFoundHandle != nil {
			res, err = exec.NotFoundHandle(ctx, ctx.CmdStr)
		} else {
			err = fmt.Errorf("%s not found", ctx.CmdStr)
		}
		return
	}
	if len(exec.params) > 0 && len(args) < len(exec.params) {
		s := ""
		for _, v := range exec.params {
			s += "{" + v + "} "
		}
		res = &Response{Code: 1000, Error: fmt.Sprintf("Usage: %s %s", ctx.CmdStr, s)}
		err = nil
	} else {
		ctx.Args = args
		ctx.params = exec.params
		if err = exec.handleFunc(ctx); err == nil {
			res = ctx.response
		} else {
			if err == ErrInvalidArgument {
				res = &Response{Code: 1000, Error: fmt.Sprint(ctx.Get("usage"))}
				err = nil
			}
		}
	}
	return
}

func NewExecutor(args ...string) *Executor {
	var (
		name        string
		usage       string
		description string
	)
	if len(args) > 0 {
		name = args[0]
	}
	if len(args) > 2 {
		description = args[2]
	}
	if len(args) > 1 {
		usage = args[1]
	}
	return &Executor{
		name:        name,
		usage:       usage,
		createdTime: time.Now(),
		description: description,
		children:    make(map[string]*Executor),
		seq:         atomic.AddInt64(&_seq, 1),
	}
}