cli/mtool/tool.go

package mtool

import (
	"text/tabwriter"
	"strings"
	"flag"
	"sort"
	"fmt"
	"io"
	"os"
)

type Tool struct {
	name string
	handler Handler
	desc, ldesc, usage string
	subs ToolMap
	parent *Tool
}

// Returns new empty tool with specified name.
func T(name string) *Tool {
	ret := &Tool{}
	ret.name = name
	return ret
}

func (t *Tool) Handler(h Handler) *Tool {
	t.handler = h
	return t
}

func (t *Tool) Func(fn HandlerFunc) *Tool {
	t.handler = fn
	return t
}

func (t *Tool) Desc(d string) *Tool {
	t.desc = d
	return t
}

func (t *Tool) Ldesc(d string) *Tool {
	t.ldesc = d
	return t
}

func (t *Tool) Usage(u string) *Tool {
	t.usage = u
	return t
}

func (t *Tool) Subs(tools ...*Tool) *Tool {
	if t.subs == nil {
		t.subs = ToolMap{}
	}
	for _, tool := range tools {
		tool.parent = t
		t.subs[tool.name] = tool
	}
	return t
}

func (t *Tool) Name() string {
	return t.name
}

func (t *Tool) FullName() string {
	ret := ""
	for t != nil {
		ret = t.name + func() string {
			if ret != "" {
				return " "
			}
			return ""
		}() + ret
		t = t.parent
	}
	return ret
}

func (t *Tool) IsRoot() bool {
	return t.parent == nil
}

func (t *Tool) ProgName() string {
	for t.parent != nil {
		t = t.parent
	}
	return t.name
}

func (t *Tool) PrintSubs(out io.Writer) {
	w := new(tabwriter.Writer)
	w.Init(out, 0, 0, 1, ' ', 0)
	defer w.Flush()
	keys := make([]string, len(t.subs))
	i := 0
	for k, _ := range t.subs {
		keys[i] = k
		i++
	}
	sort.Strings(keys)
	for _, k := range keys {
		tool := t.subs[k]
		fmt.Fprintf(w, "  %s\t%s\n", k, tool.desc)
	}
}

func (t *Tool) Run(args []string) {
	var(
		usageTool *Tool
	)
	// Should implement the shit later
	//binBase := path.Base(arg0) ;
	//binBase = binBase[:len(binBase)-len(path.Ext(binBase))]

	flagSet := flag.NewFlagSet(t.FullName(), flag.ExitOnError)
	flags := &Flags{
		FlagSet : flagSet,
		envMap: make(map[string]string),
		envNameMap: make(map[string] []string),
	}
	out := flags.Output()
	flags.Usage = func() {
		nflags := 0
		flags.VisitAll(func(f *flag.Flag){
			nflags++
			f.Usage = FormatInCode(f.Usage, false)
			varNames, ok := flags.envNameMap[f.Name]
			if !ok || len(varNames) == 0 {
				return
			}

			f.Usage += "\n("
			for i, name := range varNames {
				f.Usage += "$"+name
				if i < len(varNames) - 1 {
					f.Usage += ", "
				}
			}
			f.Usage += ")"
		})
		hasOptions := nflags != 0

		// Name
		if usageTool.desc != "" {
			fmt.Fprintf(
				out, "Name:\n  %s - %s\n\n",
				usageTool.FullName(), FormatInCode(t.desc, true),
			)
		}
		

		// Usage
		fmt.Fprintf(
			out, "Usage:\n  %s",
			usageTool.FullName(),
		)
		if hasOptions {
			fmt.Fprintf(out, " [options]")
		}
		
		if usageTool.usage != "" {
			fmt.Fprintf(
				out,
				" %s",
				usageTool.usage,
			)
		}

		fmt.Fprintln(out, "")

		if usageTool.ldesc != "" {
			fmt.Fprintf(
				out,
				"\nDescription:\n  %s\n",
				FormatInCode(usageTool.ldesc, true),
			)
		}
		
		// Options
		if hasOptions {
			fmt.Fprintln(out, "\nOptions:")
			fmt.Fprintln(out, "  -- option terminator")
			flags.PrintDefaults()
		}
		
	}
	
	flags.args = args

	// If the tool has its own handler run it.
	if t.handler != nil {
		usageTool = t
		t.handler.Handle(flags)
		return
	}

	// Print available sub commands if
	// got no arguments.
	if len(args) == 0 || func() bool {
		toolName := args[0]
		_, ok := t.subs[toolName]
		if ok {
			return false
		}
		if toolName != "" && toolName[0]=='-' {
			return true
		}
		return false
	}() {


		if t.desc != "" {
			fmt.Fprintf(
				out, "Name:\n  %s - %s\n\n",
				t.FullName(), FormatInCode(t.desc, true),
			)
		}

		fmt.Fprintf(out, "Usage:\n"+
			"  %s <command>\n", t.FullName())

		if t.ldesc != "" {
			fmt.Fprintf(
				out,
				"\nDescription:\n  %s\n",
				FormatInCode(t.ldesc, true),
			)
		}

		if len(t.subs) > 0 {
			fmt.Fprint(out, "\nCommands:\n")
			t.PrintSubs(out)
		}
		
		os.Exit(1)
	}
	toolName := args[0]
	args = args[1:]
	
	if _, ok := t.subs[toolName] ; !ok {
		fmt.Fprintf(
			out,
			"%s: No such util %q'\n" +
			"use ' %s ' to see available commands\n",
			t.FullName(), toolName,
			t.FullName(),
		)
		os.Exit(1)
	}

	sub := t.subs[toolName]
	usageTool = sub
	sub.Run(args)
}

// Returns the built-in
// string in a more printable format.
// Is used in printing descriptions.
func FormatInCode(
	desc string,
	addSpaces bool,
) string {
	if desc == "" {
		return desc
	}
	desc = strings.ReplaceAll(desc, "\t", "")
	if desc[0] == '\n' {
		desc = desc[1:]
	}
	if desc[len(desc)-1] == '\n' {
		desc = desc[:len(desc)-1]
	}
	if addSpaces {
		desc = strings.ReplaceAll(desc, "\n", "\n  ")
	}
	return desc
}