color.go

// Copyright 2016 Ben Morgan. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

// Copyright 2013, Meng Zhang. All rights reserved.
// URL: https://github.com/wsxiaoys/terminal
// File URL: https://github.com/wsxiaoys/terminal/blob/decf4e097e2e3471b254da8d30c3599d330fe7ba/color/color.go

// Package color provides printing in ANSI colors.
//
// TODO: Need to document this. See the Color function for now.
//
// This package uses code from [github.com/wsxiaoys/terminal](https://github.com/wsxiaoys/terminal). Many thanks.
package color

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"

	"golang.org/x/crypto/ssh/terminal"
)

// Mapping from character to concrete escape code.
var codeMap = map[int]int{
	'|': 0,
	'!': 1,
	'.': 2,
	'/': 3,
	'_': 4,
	'^': 5,
	'&': 6,
	'?': 7,
	'-': 8,

	'k': 30,
	'r': 31,
	'g': 32,
	'y': 33,
	'b': 34,
	'm': 35,
	'c': 36,
	'w': 37,
	'd': 39,

	'K': 40,
	'R': 41,
	'G': 42,
	'Y': 43,
	'B': 44,
	'M': 45,
	'C': 46,
	'W': 47,
	'D': 49,
}

// ErrInvalidEscape is an error that is used when the parser panics.
var ErrInvalidEscape = errors.New("invalid escape rune")
var ErrUnexpectedEOF = errors.New("unexpected EOF while parsing")

// ColorReset is the string that resets the text to default style.
const ColorReset = "\033[0m"

// ColorCode compiles a color syntax string like "rG" to escape code.
func ColorCode(s string) string {
	attr := 0
	fg := 39
	bg := 49

	for _, key := range s {
		c, ok := codeMap[int(key)]
		if !ok {
			panic("wrong color syntax: " + string(key))
		}

		switch {
		case 0 <= c && c <= 8:
			attr = c
		case 30 <= c && c <= 37:
			fg = c
		case 40 <= c && c <= 47:
			bg = c
		}
	}
	return fmt.Sprintf("\033[%d;%d;%dm", attr, fg, bg)
}

// Color translates a string into an escaped string.
//
// This example will output the text with a Blue foreground and a Black background
//      color.Println("@{bK}Example Text")
//
// This one will output the text with a red foreground
//      color.Println("@rExample Text")
//
// This one will escape the @
//      color.Println("@@")
//
// Full color syntax code
//      @{rgbcmykwRGBCMYKW}  foreground/background color
//        r/R:  Red
//        g/G:  Green
//        b/B:  Blue
//        c/C:  Cyan
//        m/M:  Magenta
//        y/Y:  Yellow
//        k/K:  Black
//        w/W:  White
//      @{|}  Reset format style
//      @{!./_} Bold / Dim / Italic / Underline
//      @{^&} Blink / Fast blink
//      @{?} Reverse the foreground and background color
//      @{-} Hide the text
// Note some of the functions are not widely supported, like "Fast blink" and "Italic".
func Color(s string, escape rune) string {
	return newParser(escape, true).translateReset(s)
}

// Decolor cleans a string of @x color codes.
func Decolor(s string, escape rune) string {
	return newParser(escape, false).translateReset(s)
}

// Uncolor cleans a string of ANSI color codes.
func Uncolor(s string) string {
	panic("not implemented")
}

type Colorizer struct {
	status string
	w      io.Writer
	*parser
}

func New() *Colorizer {
	c := &Colorizer{
		parser: newParser('@', true),
	}
	c.SetFile(os.Stdout)
	return c
}

func (c *Colorizer) EscapeChar() rune {
	return c.parser.escape
}

// String returns the current state, either always or never.
func (c Colorizer) String() string {
	return c.status
}

// Returns the name of the type: "Colorizer".
//
// TODO: This is implemented to fulfil the github.com/spf13/pflag.Value interface.
// It's not very well documented what this is for however.
func (c Colorizer) Type() string {
	return "Colorizer"
}

// Set sets c on or off, with s one of "auto", "always", or "never",
// otherwise an error is returned.
//
// Note: if "auto" is passed, and SetOutput was passed an io.Writer that
// is not an *os.File, nothing will happen.
// This lets you decide whether you want it to fall back to enabled or disabled.
func (c *Colorizer) Set(s string) (err error) {
	switch s {
	case "auto":
		if w, ok := c.w.(*os.File); ok {
			c.SetFile(w)
		}
		// TODO: Decide what to do in this case. It's a writer, but not to a file,
		// so I don't know how it will act. The conservative thing to do would be
		// to disable it, the liberal thing would be to allow the user (programmer)
		// to decide before-hand what will happen.
	case "always":
		c.SetEnabled(true)
	case "never":
		c.SetEnabled(false)
	default:
		err = errors.New("expect one of auto, always, or never")
	}
	return err
}

// SetEscapeChar sets the escape character, which can be one of the following characters:
//
//		@ * + = ~
//
// If it is none of these characters, then this function panics with ErrInvalidEscape.
func (c *Colorizer) SetEscapeChar(r rune) {
	if c.EscapeChar() == r {
		return
	}

	for _, q := range []rune{'*', '@', '+', '=', '~'} {
		if r == q {
			c.parser.escape = r
			return
		}
	}

	panic(ErrInvalidEscape)
}

func (c *Colorizer) Enabled() bool {
	return c.parser.color
}

func (c *Colorizer) SetEnabled(b bool) {
	if b {
		c.status = "always"
	} else {
		c.status = "never"
	}
	c.parser.color = b
}

func (c *Colorizer) SetOutput(w io.Writer) {
	// TODO: should I set anything here?
	c.w = w
}

func (c *Colorizer) SetFile(f *os.File) {
	c.SetEnabled(terminal.IsTerminal(int(f.Fd())))
	c.status = "auto"
	c.w = f
}

func (c *Colorizer) Color(s string) string {
	return c.translateReset(s)
}

func (c *Colorizer) colorAny(args []interface{}) []interface{} {
	n := len(args)
	r := make([]interface{}, n, n+1)
	for i, x := range args {
		if str, ok := x.(string); ok {
			x = c.translateOnly(str)
		}
		r[i] = x
	}
	if c.Enabled() {
		r = append(r, ColorReset)
	}
	return r
}

func (c *Colorizer) Print(a ...interface{}) (int, error) {
	return fmt.Fprint(c.w, c.colorAny(a)...)
}
func (c *Colorizer) Println(a ...interface{}) (int, error) {
	return fmt.Fprintln(c.w, c.colorAny(a)...)
}
func (c *Colorizer) Printf(format string, a ...interface{}) (int, error) {
	return fmt.Fprintf(c.w, c.translateReset(format), a...)
}
func (c *Colorizer) Fprint(w io.Writer, a ...interface{}) (int, error) {
	return fmt.Fprint(w, c.colorAny(a)...)
}
func (c *Colorizer) Fprintln(w io.Writer, a ...interface{}) (int, error) {
	return fmt.Fprintln(w, c.colorAny(a)...)
}
func (c *Colorizer) Fprintf(w io.Writer, format string, a ...interface{}) (int, error) {
	return fmt.Fprintf(w, c.translateReset(format), a...)
}
func (c *Colorizer) Sprint(a ...interface{}) string {
	return fmt.Sprint(c.colorAny(a)...)
}
func (c *Colorizer) Sprintln(a ...interface{}) string {
	return fmt.Sprintln(c.colorAny(a)...)
}
func (c *Colorizer) Sprintf(format string, a ...interface{}) string {
	return fmt.Sprintf(c.translateReset(format), a...)
}

type parser struct {
	escape rune
	color  bool
}

func newParser(escape rune, color bool) *parser {
	return &parser{
		escape: escape,
		color:  color,
	}
}

type handler func(p *parser, in, out *bytes.Buffer) (handler, error)

func (p *parser) translateReset(s string) string {
	in := bytes.NewBufferString(s)
	out := bytes.NewBufferString("")

	var h = handleRegular
	var err error
	for {
		h, err = h(p, in, out)
		if err != nil {
			panic(err)
		}
		if h == nil {
			break
		}
	}
	if p.color {
		out.WriteString(ColorReset)
	}
	return out.String()
}

func (p *parser) translateOnly(s string) string {
	in := bytes.NewBufferString(s)
	out := bytes.NewBufferString("")

	var h = handleRegular
	var err error
	for {
		h, err = h(p, in, out)
		if err != nil {
			panic(err)
		}
		if h == nil {
			break
		}
	}
	return out.String()
}

func handleRegular(p *parser, in, out *bytes.Buffer) (handler, error) {
	for {
		r, _, err := in.ReadRune()
		// The only error that can happen here is that we have reached the end of file,
		// or that a rune is messed up. If the rune is messed up, we treat it normally.
		// This is why we only check for io.EOF.
		if err == io.EOF {
			break
		}

		if r == p.escape {
			return handleEscape, nil
		}
		out.WriteRune(r)
	}
	return nil, nil
}

func handleEscape(p *parser, in, out *bytes.Buffer) (handler, error) {
	r, _, err := in.ReadRune()
	if err == io.EOF {
		return nil, ErrUnexpectedEOF
	}

	if r == '{' {
		return handleEscapeClause, nil
	} else if r == p.escape {
		out.WriteRune(p.escape)
	} else if p.color {
		out.WriteString(ColorCode(string(r)))
	}
	return handleRegular, nil
}

func handleEscapeClause(p *parser, in, out *bytes.Buffer) (handler, error) {
	bs := bytes.NewBufferString("")
	for {
		r, _, err := in.ReadRune()
		if err == io.EOF {
			return nil, ErrUnexpectedEOF
		}

		if r == '}' {
			break
		}
		bs.WriteRune(r)
	}

	if p.color {
		out.WriteString(ColorCode(bs.String()))
	}
	return handleRegular, nil
}