Add color support for multi-line diffs

Author: rliebz

be/assertions.go

@@ -6,8 +6,6 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/google/go-cmp/cmp"
-
 	"github.com/rliebz/ghost"
 	"github.com/rliebz/ghost/ghostlib"
 	"github.com/rliebz/ghost/internal/constraints"
@@ -114,14 +112,10 @@ func DeepEqual[T any](got, want T) ghost.Result {
 	args := ghostlib.ArgsFromAST(got, want)
 	argGot, argWant := args[0], args[1]
 
-	if diff := cmp.Diff(
-		want, got,
-		cmp.Exporter(func(reflect.Type) bool { return true }),
-	); diff != "" {
+	if diff := colorDiff(want, got); diff != "" {
 		return ghost.Result{
 			Ok: false,
 			Message: fmt.Sprintf(`%v != %v
-diff (-want +got):
 %v`, argGot, argWant, diff),
 		}
 	}
@@ -169,8 +163,7 @@ value: %v
 		return ghost.Result{
 			Ok: false,
 			Message: fmt.Sprintf(`%v != %v
-diff (-want +got):
-%v`, argGot, argWant, cmp.Diff(want, got)),
+%v`, argGot, argWant, colorDiff(want, got)),
 		}
 	case reflect.String:
 		if strings.ContainsAny(v.String(), "\n\r") ||
@@ -179,8 +172,7 @@ diff (-want +got):
 			return ghost.Result{
 				Ok: false,
 				Message: fmt.Sprintf(`%v != %v
-diff (-want +got):
-%v`, argGot, argWant, cmp.Diff(want, got)),
+%v`, argGot, argWant, colorDiff(want, got)),
 			}
 		}
 
@@ -235,7 +227,7 @@ func JSONEqual[T ~string | ~[]byte](got, want T) ghost.Result {
 	args := ghostlib.ArgsFromAST(got, want)
 	argGot, argWant := args[0], args[1]
 
-	diff, kind := jsondiff.Diff(got, want)
+	diff, kind := colorJSONDiff(got, want)
 
 	switch kind {
 	case jsondiff.Match:
@@ -270,7 +262,7 @@ want:
 	return ghost.Result{
 		Ok: false,
 		Message: fmt.Sprintf(`%v and %v are not JSON equal
-diff: %s`, argGot, argWant, diff),
+%s`, argGot, argWant, diff),
 	}
 }
 

be/diff.go

@@ -0,0 +1,51 @@
+package be
+
+import (
+	"fmt"
+	"reflect"
+	"strings"
+
+	"github.com/google/go-cmp/cmp"
+
+	"github.com/rliebz/ghost/internal/color"
+	"github.com/rliebz/ghost/internal/jsondiff"
+)
+
+var exportTypes = cmp.Exporter(func(reflect.Type) bool { return true })
+
+func colorDiff[T any](x, y T, opts ...cmp.Option) string {
+	diff := cmp.Diff(x, y, append(opts, exportTypes)...)
+	return applyColors(diff)
+}
+
+func colorJSONDiff[T ~string | ~[]byte](got, want T) (string, jsondiff.Kind) {
+	diff, kind := jsondiff.Diff(got, want)
+	return applyColors(diff), kind
+}
+
+func applyColors(diff string) string {
+	if diff == "" {
+		return ""
+	}
+
+	ss := strings.Split(diff, "\n")
+	for i, s := range ss {
+		switch {
+		case strings.HasPrefix(s, "-"):
+			ss[i] = color.Red(s)
+		case strings.HasPrefix(s, "+"):
+			ss[i] = color.Green(s)
+		// Only color the first character, since we expect inline red/green
+		case strings.HasPrefix(s, "~"):
+			ss[i] = color.Yellow("~") + s[1:]
+		}
+	}
+
+	return fmt.Sprintf(
+		`diff (%s %s):
+%v`,
+		color.Red("-want"),
+		color.Green("+got"),
+		strings.Join(ss, "\n"),
+	)
+}

be/main_test.go

@@ -0,0 +1,14 @@
+package be
+
+import (
+	"os"
+	"testing"
+)
+
+// Avoid dealing with ANSI escape sequences.
+func TestMain(m *testing.M) {
+	if err := os.Setenv("NO_COLOR", "1"); err != nil {
+		panic(err)
+	}
+	os.Exit(m.Run())
+}

internal/color/color.go

@@ -0,0 +1,66 @@
+// Package color exports utilities for applying color to terminal outputs.
+package color
+
+import (
+	"os"
+	"regexp"
+	"runtime"
+	"sync"
+)
+
+const (
+	// ANSIReset is the ANSI escape sequence for reset.
+	ANSIReset = "\033[0m"
+	// ANSIRed is the ANSI escape sequence for red.
+	ANSIRed = "\033[31m"
+	// ANSIGreen is the ANSI escape sequence for green.
+	ANSIGreen = "\033[32m"
+	// ANSIYellow is the ANSI escape sequence for yellow.
+	ANSIYellow = "\033[33m"
+)
+
+// Red colors the string red.
+func Red(s string) string {
+	return apply(ANSIRed, s)
+}
+
+// Green colors the string green.
+func Green(s string) string {
+	return apply(ANSIGreen, s)
+}
+
+// Yellow colors the string yellow
+func Yellow(s string) string {
+	return apply(ANSIYellow, s)
+}
+
+var (
+	// reReset captures any reset sequence that is followed by at least one
+	// character. The extra character check prevents us from writing escape
+	// squences at the end of lines.
+	reReset = regexp.MustCompile(`(\033\[0m)(.)`)
+	// reANSI identifies any non-reset ANSI escape sequence.
+	reANSI = regexp.MustCompile(`\033\[[1-9][\d;]*m`)
+)
+
+func apply(color string, s string) string {
+	if !Enabled() {
+		return s
+	}
+
+	// Preserve existing colors by resetting before ANSI escape sequences
+	// and re-applying after the reset sequence.
+	s = reReset.ReplaceAllString(s, "$1"+color+"$2")
+	s = reANSI.ReplaceAllString(s, ANSIReset+"$0")
+	return color + s + ANSIReset
+}
+
+// Enabled returns whether colors are enabled.
+var Enabled = sync.OnceValue(func() bool {
+	if runtime.GOOS == "windows" {
+		return false
+	}
+
+	_, ok := os.LookupEnv("NO_COLOR")
+	return !ok
+})

internal/jsondiff/jsondiff.go

@@ -9,6 +9,8 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+
+	"github.com/rliebz/ghost/internal/color"
 )
 
 // Kind describes the result of the diff categorically.
@@ -327,12 +329,25 @@ func (d *differ) writeMismatch(got, want any) {
 		d.buf.Write(data)
 	}
 
+	d.writeANSI(color.ANSIRed)
 	d.writeValueInline(want)
+	d.writeANSI(color.ANSIReset)
+
 	d.buf.WriteString(" => ")
+
+	d.writeANSI(color.ANSIGreen)
 	d.writeValueInline(got)
+	d.writeANSI(color.ANSIReset)
+
 	d.failMatch()
 }
 
+func (d *differ) writeANSI(sequence string) {
+	if color.Enabled() {
+		d.buf.WriteString(sequence)
+	}
+}
+
 // TODO: I would rather have the last element of EACH collection handle the
 // middle commas correctly. Probably a hefty refactor, can save it for later.
 func (d *differ) writeElementSeparator(first bool) {

internal/jsondiff/jsondiff_test.go

@@ -1,13 +1,22 @@
 package jsondiff_test
 
 import (
+	"os"
 	"testing"
 
 	"github.com/rliebz/ghost"
 	"github.com/rliebz/ghost/be"
 	"github.com/rliebz/ghost/internal/jsondiff"
 )
 
+// Avoid dealing with ANSI escape sequences in tests.
+func TestMain(m *testing.M) {
+	if err := os.Setenv("NO_COLOR", "1"); err != nil {
+		panic(err)
+	}
+	os.Exit(m.Run())
+}
+
 func TestDiff(t *testing.T) {
 	tests := []struct {
 		name     string