m

kama.go

@@ -12,6 +12,7 @@ package kama
 
 import (
 	"fmt"
+	"os"
 	"reflect"
 	"strings"
 	"sync"
@@ -133,7 +134,7 @@ func Dir(i interface{}, c ...Config) string {
 // Stack trace from the runtime/stdlib is colored blue, third party libraries is yellow
 // whereas your code is red.
 func Stackp() {
-	stackp()
+	stackp(os.Stderr)
 }
 
 // Diffp prints a formatted diff showing the minimum line-level additions and removals that would turn old into new.

stack.go

@@ -3,6 +3,7 @@ package kama
 import (
 	"bufio"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -16,13 +17,14 @@ const (
 	yourColor       = "red"
 )
 
-func stackp() {
+func stackp(w io.Writer) {
 	goModCache := os.Getenv("GOMODCACHE")
 	re := regexp.MustCompile(`\d:`) // this pattern is the one created in `readLastLine()`
 
 	traces := getStackTrace()
 	if len(traces) > 0 {
 		printWithColor(
+			w,
 			fmt.Sprintf("LEGEND:\n compiler: %s\n thirdParty: %s\n yours: %s\n", runtimeColor, thirdPartyColor, yourColor),
 			"DEFAULT",
 			true,
@@ -32,16 +34,16 @@ func stackp() {
 	for _, v := range traces {
 		if strings.Contains(v, "go/src/") {
 			// compiler
-			printWithColor(v, runtimeColor, false)
+			printWithColor(w, v, runtimeColor, false)
 		} else if goModCache != "" && strings.Contains(v, goModCache) {
 			// third party
-			printWithColor(v, thirdPartyColor, false)
+			printWithColor(w, v, thirdPartyColor, false)
 		} else if re.MatchString(v) {
 			// this is code snippets
-			printWithColor(v, yourColor, false)
+			printWithColor(w, v, yourColor, false)
 		} else {
 			// your code
-			printWithColor(v, yourColor, true)
+			printWithColor(w, v, yourColor, true)
 		}
 	}
 }
@@ -142,22 +144,22 @@ func readLastLine(file string, line int64) string {
 	return strings.TrimSuffix(txt, "\n")
 }
 
-func reset() {
+func reset(w io.Writer) {
 	const escape = "\x1b"
 	const r = 0
-	fmt.Fprintf(os.Stderr, "%s[%dm", escape, r)
+	fmt.Fprintf(w, "%s[%dm", escape, r)
 }
 
-func setColor(code int, bold bool) {
+func setColor(w io.Writer, code int, bold bool) {
 	const escape = "\x1b"
 	if bold {
-		fmt.Fprintf(os.Stderr, "%s[1%dm", escape, code)
+		fmt.Fprintf(w, "%s[1%dm", escape, code)
 	} else {
-		fmt.Fprintf(os.Stderr, "%s[%dm", escape, code)
+		fmt.Fprintf(w, "%s[%dm", escape, code)
 	}
 }
 
-func printWithColor(s, color string, bold bool) {
+func printWithColor(w io.Writer, s, color string, bold bool) {
 	// TODO: should be iota
 	colors := map[string]int{
 		// Go compiler == compilerColor
@@ -178,12 +180,12 @@ func printWithColor(s, color string, bold bool) {
 	}
 
 	if noColor() {
-		fmt.Fprintln(os.Stderr, s)
+		fmt.Fprintln(w, s)
 	} else {
-		defer reset()
+		defer reset(w)
 		color = strings.ToLower(color)
-		setColor(colors[color], bold)
-		fmt.Fprintln(os.Stderr, s)
+		setColor(w, colors[color], bold)
+		fmt.Fprintln(w, s)
 	}
 }
 

stack_test.go

@@ -1,6 +1,8 @@
 package kama
 
 import (
+	"bytes"
+	"io"
 	"testing"
 
 	"go.akshayshah.org/attest"
@@ -37,10 +39,12 @@ func Test_stackp(t *testing.T) {
 	t.Run("test-stackp", func(t *testing.T) {
 		t.Parallel()
 
-		d()
+		w := &bytes.Buffer{}
+
+		d(w)
 	})
 }
 
-func d() {
-	stackp()
+func d(w io.Writer) {
+	stackp(w)
 }