Color modes

main.go

@@ -1,12 +1,15 @@
 package main
 
 import (
-	"bufio"
+	"bytes"
+	"flag"
 	"fmt"
 	"io"
 	"math"
 	"os"
 	"os/exec"
+	"sync"
+	"unicode/utf8"
 )
 
 func rbg(i int) (int, int, int) {
@@ -16,43 +19,176 @@ func rbg(i int) (int, int, int) {
 		int(math.Sin(f*float64(i)+4*math.Pi/3)*127 + 128)
 }
 
-func color_printer(output []rune) {
-	for i := 0; i < len(output); i++ {
-		r, g, b := rbg(i)
-		fmt.Printf("\033[38;2;%d;%d;%dm%c\033[0m", r, g, b, output[i])
+func fire(i int) (int, int, int) {
+	if i%3 == 0 {
+		return 255, 0, 0
+	} else if i%3 == 1 {
+		return 255, 165, 0
 	}
-	fmt.Println()
+	return 255, 255, 0
 }
 
-func main() {
-	var output []rune
-	var cmd1 *exec.Cmd
+func ocean(i int) (int, int, int) {
+	if i%3 == 0 {
+		return 0, 255, 255
+	} else if i%3 == 1 {
+		return 0, 0, 255
+	}
+	return 75, 0, 130
+}
+
+func forest(i int) (int, int, int) {
+	if i%3 == 0 {
+		return 0, 255, 0
+	} else if i%3 == 1 {
+		return 34, 139, 34
+	}
+	return 0, 100, 0
+}
 
-	cli_arg := os.Args
-	cli_args_etc := cli_arg[2:]
-	if len(cli_args_etc) == 0 {
-		cmd1 = exec.Command(cli_arg[1])
-	} else {
-		cmd1 = exec.Command(cli_arg[1], cli_args_etc...)
+func sunset(i int) (int, int, int) {
+	if i%3 == 0 {
+		return 255, 140, 0
+	} else if i%3 == 1 {
+		return 255, 105, 180
 	}
+	return 186, 85, 211
+}
 
-	cmd1_op, _ := cmd1.StdoutPipe()
+type rainbowWriter struct {
+	dst       io.Writer
+	colorIdx  int
+	colorMode string
+	colorFunc func(int) (int, int, int)
+}
 
-	if err := cmd1.Start(); err != nil {
-		return
+func newRainbowWriter(w io.Writer, mode string) *rainbowWriter {
+	var colorFunc func(int) (int, int, int)
+
+	switch mode {
+	case "fire":
+		colorFunc = fire
+	case "ocean":
+		colorFunc = ocean
+	case "forest":
+		colorFunc = forest
+	case "sunset":
+		colorFunc = sunset
+	default:
+		colorFunc = rbg
 	}
 
-	reader := bufio.NewReader(cmd1_op)
-	for {
-		input, _, err := reader.ReadRune()
-		if err != nil && err == io.EOF {
-			break
+	return &rainbowWriter{dst: w, colorMode: mode, colorFunc: colorFunc}
+}
+
+func (w *rainbowWriter) stripANSI(data []byte) []byte {
+	out := make([]byte, 0, len(data))
+	i := 0
+	for i < len(data) {
+		if data[i] == 0x1b && i+1 < len(data) && data[i+1] == '[' {
+			i += 2
+			for i < len(data) && (data[i] < 0x40 || data[i] > 0x7e) {
+				i++
+			}
+			if i < len(data) {
+				i++
+			}
+		} else {
+			out = append(out, data[i])
+			i++
 		}
-		output = append(output, input)
 	}
-	// println("gotokekll")
+	return out
+}
+
+func (w *rainbowWriter) rainbowize(data []byte) []byte {
+	out := make([]byte, 0, len(data)*20)
+	colorIdx := w.colorIdx
+
+	for len(data) > 0 {
+		r, size := utf8.DecodeRune(data)
+		if r == utf8.RuneError && size == 1 {
+			out = append(out, data[0])
+			data = data[1:]
+			continue
+		}
 
-	cmd1.Wait()
+		switch r {
+		case '\n', '\r', '\t', ' ':
+			out = append(out, data[:size]...)
+		default:
+			r, g, b := w.colorFunc(colorIdx)
+			out = append(out, []byte(fmt.Sprintf("\033[38;2;%d;%d;%dm%s\033[0m", r, g, b, string(r)))...)
+			colorIdx++
+		}
+		data = data[size:]
+	}
+
+	w.colorIdx = colorIdx
+	return out
+}
 
-	color_printer(output)
+func (w *rainbowWriter) Write(p []byte) (int, error) {
+	stripped := w.stripANSI(p)
+	colored := w.rainbowize(stripped)
+	_, err := w.dst.Write(colored)
+	return len(p), err
 }
+
+func main() {
+	colorMode := flag.String("c", "rbg", "Color mode: rbg, fire, ocean, forest, sunset")
+	flag.Parse()
+
+	args := flag.Args()
+	if len(args) < 1 {
+		fmt.Fprintln(os.Stderr, "usage: gorainbow [-c colormode] <command> [args...]")
+		fmt.Fprintln(os.Stderr, "color modes: rbg (default), fire, ocean, forest, sunset")
+		os.Exit(2)
+	}
+
+	cmd := exec.Command(args[0], args[1:]...)
+
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "stdout pipe:", err)
+		os.Exit(1)
+	}
+
+	stderr, err := cmd.StderrPipe()
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "stderr pipe:", err)
+		os.Exit(1)
+	}
+
+	if err := cmd.Start(); err != nil {
+		fmt.Fprintln(os.Stderr, "start:", err)
+		os.Exit(1)
+	}
+
+	rwOut := newRainbowWriter(os.Stdout, *colorMode)
+	rwErr := newRainbowWriter(os.Stderr, *colorMode)
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		io.Copy(rwOut, stdout)
+	}()
+
+	go func() {
+		defer wg.Done()
+		io.Copy(rwErr, stderr)
+	}()
+
+	waitErr := cmd.Wait()
+	wg.Wait()
+
+	if waitErr != nil {
+		if ee, ok := waitErr.(*exec.ExitError); ok {
+			os.Exit(ee.ExitCode())
+		}
+		fmt.Fprintln(os.Stderr, waitErr)
+		os.Exit(1)
+	}
+}