Improve fatal and color packages (#4)

* feat!: Add ability to register func to run before exiting in fatal

BREAKING CHANGE: fatal.ShowStackTraces is now a function

* feat: Improve color methods to strip reset values and add tests

* Oops

color/color.go

@@ -1,41 +1,75 @@
 package color
 
+import (
+	"fmt"
+	"os"
+	"regexp"
+)
+
 const (
-	red     = "\x1b[31m"
-	green   = "\x1b[32m"
-	yellow  = "\x1b[33m"
-	blue    = "\x1b[34m"
-	magenta = "\x1b[35m"
-	cyan    = "\x1b[36m"
-	reset   = "\x1b[0m"
+	ansiFgRed     = 31
+	ansiFgGreen   = 32
+	ansiFgYellow  = 33
+	ansiFgBlue    = 34
+	ansiFgMagenta = 35
+	ansiFgCyan    = 36
+	asnsiFgWhite  = 37
+	ansiResetFg   = 39
 )
 
+// Support for NO_COLOR env var
+// https://no-color.org/
+var noColor = false
+
+func init() {
+	// The standard says the value doesn't matter, only whether or not it's set
+	if _, ok := os.LookupEnv("NO_COLOR"); ok {
+		noColor = true
+	}
+}
+
+func apply(str string, start, end int) string {
+	if noColor {
+		return str
+	}
+
+	regex := regexp.MustCompile(fmt.Sprintf("\\x1b\\[%dm", end))
+	// Remove any occurrences of reset to make sure color isn't messed up
+	sanitized := regex.ReplaceAllString(str, "")
+	return fmt.Sprintf("\x1b[%dm%s\x1b[%dm", start, sanitized, end)
+}
+
 // Red creates a red colored string
 func Red(str string) string {
-	return red + str + reset
+	return apply(str, ansiFgRed, ansiResetFg)
 }
 
 // Green creates a green colored string
 func Green(str string) string {
-	return green + str + reset
+	return apply(str, ansiFgGreen, ansiResetFg)
 }
 
 // Yellow creates a yellow colored string
 func Yellow(str string) string {
-	return yellow + str + reset
+	return apply(str, ansiFgYellow, ansiResetFg)
 }
 
 // Blue creates a blue colored string
 func Blue(str string) string {
-	return blue + str + reset
+	return apply(str, ansiFgBlue, ansiResetFg)
 }
 
 // Magenta creates a magenta colored string
 func Magenta(str string) string {
-	return magenta + str + reset
+	return apply(str, ansiFgMagenta, ansiResetFg)
 }
 
 // Cyan creates a cyan colored string
 func Cyan(str string) string {
-	return cyan + str + reset
+	return apply(str, ansiFgCyan, ansiResetFg)
+}
+
+// White creates a white colored string
+func White(str string) string {
+	return apply(str, asnsiFgWhite, ansiResetFg)
 }

color/color_test.go

@@ -0,0 +1,82 @@
+package color
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestColors(t *testing.T) {
+	noColor = false
+
+	tests := []struct {
+		name     string
+		colorFn  func(string) string
+		input    string
+		expected string
+	}{
+		{
+			"Red() test",
+			Red,
+			"foo bar",
+			"\x1b[31mfoo bar\x1b[39m",
+		},
+		{
+			"Green() test",
+			Green,
+			"foo bar",
+			"\x1b[32mfoo bar\x1b[39m",
+		},
+		{
+			"Yellow() test",
+			Yellow,
+			"foo bar",
+			"\x1b[33mfoo bar\x1b[39m",
+		},
+		{
+			"Blue() test",
+			Blue,
+			"foo bar",
+			"\x1b[34mfoo bar\x1b[39m",
+		},
+		{
+			"Magenta() test",
+			Magenta,
+			"foo bar",
+			"\x1b[35mfoo bar\x1b[39m",
+		},
+		{
+			"Cyan() test",
+			Cyan,
+			"foo bar",
+			"\x1b[36mfoo bar\x1b[39m",
+		},
+		{
+			"White() test",
+			White,
+			"foo bar",
+			"\x1b[37mfoo bar\x1b[39m",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			received := tt.colorFn(tt.input)
+			assert.Equal(t, tt.expected, received)
+		})
+	}
+}
+
+func TestStripReset(t *testing.T) {
+	noColor = false
+
+	received := Red("foo \x1b[39mbar")
+	assert.Equal(t, "\x1b[31mfoo bar\x1b[39m", received)
+}
+
+func TestNoColor(t *testing.T) {
+	noColor = true
+
+	received := Red("foo bar")
+	assert.Equal(t, "foo bar", received)
+}

fatal/fatal.go

@@ -2,43 +2,84 @@ package fatal
 
 import (
 	"fmt"
+	"io"
 	"os"
 )
 
-var ShowStackTraces = true
+// Package state
+var (
+	shouldShowStackTraces = false
+	onExitHandler         func()
+)
+
+// Used for dependency injection in tests
+// Normally having tests touch private stuff is bad
+// but this is the only way I could figure out to mock os.Exit
+var (
+	errWriter io.Writer      = os.Stderr
+	exitFunc  func(code int) = os.Exit
+)
 
+// ShowStackTraces sets whether or not stack traces should be printed
+// when ExitErr and ExitErrf are called.
+func ShowStackTraces(show bool) {
+	shouldShowStackTraces = show
+}
+
+// OnExit registers a handler that will run before os.Exit is called.
+// This is useful for performing any clean up that would usually be called
+// in a defer block since defers are not called when os.Exit is used.
+func OnExit(handler func()) {
+	onExitHandler = handler
+}
+
+// ExitErr prints the given message and error to stderr then exits the program.
 func ExitErr(err error, message string) {
-	fmt.Fprintln(os.Stderr, message)
+	fmt.Fprintln(errWriter, message)
+
+	if err != nil {
+		if shouldShowStackTraces {
+			fmt.Fprintf(errWriter, "Error: %+v\n", err)
+		} else {
+			fmt.Fprintf(errWriter, "Error: %s\n", err)
+		}
+	}
 
-	if ShowStackTraces {
-		fmt.Fprintf(os.Stderr, "Error: %+v\n", err)
-	} else {
-		fmt.Fprintf(os.Stderr, "Error: %s\n", err)
+	if onExitHandler != nil {
+		onExitHandler()
 	}
 
-	os.Exit(1)
+	exitFunc(1)
 }
 
+// ExitErrf prints the given message and error to stderr then exits the program.
+// Supports printf like formatting.
 func ExitErrf(err error, format string, a ...interface{}) {
-	fmt.Fprintf(os.Stderr, format, a...)
-	fmt.Fprintln(os.Stderr)
+	fmt.Fprintf(errWriter, format, a...)
+	fmt.Fprintln(errWriter)
+
+	if err != nil {
+		if shouldShowStackTraces {
+			fmt.Fprintf(errWriter, "Error: %+v\n", err)
+		} else {
+			fmt.Fprintf(errWriter, "Error: %s\n", err)
+		}
+	}
 
-	if ShowStackTraces {
-		fmt.Fprintf(os.Stderr, "Error: %+v\n", err)
-	} else {
-		fmt.Fprintf(os.Stderr, "Error: %s\n", err)
+	if onExitHandler != nil {
+		onExitHandler()
 	}
 
-	os.Exit(1)
+	exitFunc(1)
 }
 
+// Exit prints the given message to stderr then exists the program.
 func Exit(message string) {
-	fmt.Fprintln(os.Stderr, message)
-	os.Exit(1)
+	ExitErr(nil, message)
 }
 
+// Exitf prints the given message to stderr then exits the program.
+// Supports printf like formatting.
 func Exitf(format string, a ...interface{}) {
-	fmt.Fprintf(os.Stderr, format, a...)
-	fmt.Fprintln(os.Stderr)
-	os.Exit(1)
+	ExitErrf(nil, format, a...)
 }