kitty/tools/cli/infrastructure.go

package cli

import (
	"fmt"
	"os"
	"strings"
	"syscall"
	"unicode"
	"unsafe"

	"github.com/fatih/color"
	runewidth "github.com/mattn/go-runewidth"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"

	"kitty"
)

type Winsize struct {
	Rows    uint16
	Cols    uint16
	Xpixels uint16
	Ypixels uint16
}

func GetTTYSize() (Winsize, error) {
	var ws Winsize
	f, err := os.OpenFile("/dev/tty", os.O_RDONLY, 0755)
	if err != nil {
		return ws, err
	}
	fd := f.Fd()
	retCode, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), uintptr(syscall.TIOCGWINSZ), uintptr(unsafe.Pointer(&ws)))
	f.Close()

	if int(retCode) == -1 {
		return ws, errno
	}
	return ws, nil
}

func add_choices(cmd *cobra.Command, flags *pflag.FlagSet, choices []string, name string, usage string) {
	flags.String(name, choices[0], usage)
	cmd.Annotations["choices-"+name] = strings.Join(choices, "\000")
}

func Choices(cmd *cobra.Command, name string, usage string, choices ...string) {
	add_choices(cmd, cmd.Flags(), choices, name, usage)
}

func PersistentChoices(cmd *cobra.Command, name string, usage string, choices ...string) {
	add_choices(cmd, cmd.PersistentFlags(), choices, name, usage)
}

func key_in_slice(vals []string, key string) bool {
	for _, q := range vals {
		if q == key {
			return true
		}
	}
	return false
}

func ValidateChoices(cmd *cobra.Command, args []string) error {
	for key, val := range cmd.Annotations {
		if strings.HasPrefix(key, "choices-") {
			allowed := strings.Split(val, "\000")
			name := key[len("choices-"):]
			if cval, err := cmd.Flags().GetString(name); err == nil && !key_in_slice(allowed, cval) {
				return fmt.Errorf("%s: Invalid value: %s. Allowed values are: %s", color.YellowString("--"+name), color.RedString(cval), strings.Join(allowed, ", "))
			}
		}
	}
	return nil
}

var title_fmt = color.New(color.FgBlue, color.Bold).SprintFunc()
var exe_fmt = color.New(color.FgYellow, color.Bold).SprintFunc()
var opt_fmt = color.New(color.FgGreen).SprintFunc()
var italic_fmt = color.New(color.Italic).SprintFunc()

func cmd_name(cmd *cobra.Command) string {
	if cmd.Annotations != nil {
		parts := strings.Split(cmd.Annotations["exe"], " ")
		return parts[len(parts)-1]
	}
	return cmd.Name()
}

func print_created_by(root *cobra.Command) {
	fmt.Println(italic_fmt(root.Annotations["exe"]), opt_fmt(root.Version), "created by", title_fmt("Kovid Goyal"))
}

func print_with_indent(text string, indent string, screen_width int) {
	x := len(indent)
	fmt.Print(indent)
	in_sgr := false
	current_word := ""

	print_word := func(r rune) {
		w := runewidth.StringWidth(current_word)
		if x+w > screen_width {
			fmt.Println()
			fmt.Print(indent)
			x = len(indent)
			current_word = strings.TrimSpace(current_word)
		}
		fmt.Print(current_word)
		current_word = string(r)
		x += w
	}

	for _, r := range text {
		if in_sgr {
			if r == 'm' {
				in_sgr = false
			}
			continue
		}
		if r == 0x1b {
			in_sgr = true
			continue
		}
		if current_word != "" && unicode.IsSpace(r) && r != 0xa0 {
			print_word(r)
		} else {
			current_word += string(r)
		}
	}
	if current_word != "" {
		print_word(' ')
		fmt.Print(current_word)
	}
	if len(text) > 0 {
		fmt.Println()
	}
}

func show_usage(cmd *cobra.Command) error {
	ws, tty_size_err := GetTTYSize()
	screen_width := 80
	if tty_size_err == nil && ws.Cols < 80 {
		screen_width = int(ws.Cols)
	}
	fmt.Println(title_fmt("Usage")+":", exe_fmt(cmd.Annotations["exe"]), cmd.Use)
	fmt.Println()
	if len(cmd.Long) > 0 {
		print_with_indent(cmd.Long, "", screen_width)
	} else if len(cmd.Short) > 0 {
		print_with_indent(cmd.Short, "", screen_width)
	}
	if cmd.HasAvailableSubCommands() {
		fmt.Println()
		fmt.Println(title_fmt("Commands") + ":")
		for _, child := range cmd.Commands() {
			fmt.Println(" ", opt_fmt(cmd_name(child)))
			print_with_indent(child.Short, "    ", screen_width)
		}
		fmt.Println()
		print_with_indent("Get help for an individual command by running:", "", screen_width)
		fmt.Println("   ", cmd.Annotations["exe"], italic_fmt("command"), "-h")
	}
	if cmd.HasAvailableFlags() {
		options_title := cmd.Annotations["options_title"]
		if len(options_title) == 0 {
			options_title = "Options"
		}
		fmt.Println()
		fmt.Println(title_fmt(options_title) + ":")
		flag_set := cmd.LocalFlags()
		flag_set.VisitAll(func(flag *pflag.Flag) {
			fmt.Print(opt_fmt("  --" + flag.Name))
			if flag.Shorthand != "" {
				fmt.Print(", ", opt_fmt("-"+flag.Shorthand))
			}
			defval := ""
			switch flag.Value.Type() {
			default:
				defval = fmt.Sprintf("[=%s]", italic_fmt(flag.DefValue))
			case "bool":
			case "count":
			}
			if defval != "" {
				fmt.Print(" ", defval)
			}
			fmt.Println()
			if flag.Name == "help" {
				fmt.Println("   ", "Print this help message")
				fmt.Println()
				return
			}
			for _, line := range strings.Split(flag.Usage, "\n") {
				print_with_indent(line, "    ", screen_width)
			}
			if cmd.Annotations["choices-"+flag.Name] != "" {
				fmt.Println("    Choices:", strings.Join(strings.Split(cmd.Annotations["choices-"+flag.Name], "\000"), ", "))
			}
			fmt.Println()
		})
	}
	if cmd.HasParent() {
		cmd.VisitParents(func(cmd *cobra.Command) {
			if !cmd.HasParent() {
				print_created_by(cmd)
			}
		})
	} else {
		print_created_by(cmd)
	}

	return nil
}

func CreateCommand(cmd *cobra.Command, exe string) *cobra.Command {
	cmd.Annotations = make(map[string]string)
	cmd.Annotations["exe"] = exe
	return cmd
}

func Init(root *cobra.Command) {
	root.Version = kitty.VersionString
	root.PersistentPreRunE = ValidateChoices
	root.SetUsageFunc(show_usage)
	root.SetHelpTemplate("{{.UsageString}}")
}