From 02313503e64788118a33720724751046d5eec889 Mon Sep 17 00:00:00 2001
From: Tom Fleet <tomfleet2018@gmail.com>
Date: Sun, 5 Jan 2025 15:48:48 +0000
Subject: [PATCH] Support `$NO_COLOR` and `$FORCE_COLOR` (#48)

---
 internal/colour/colour.go | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/internal/colour/colour.go b/internal/colour/colour.go
index aa63ae1..09818db 100644
--- a/internal/colour/colour.go
+++ b/internal/colour/colour.go
@@ -1,6 +1,11 @@
 // Package colour implements basic text colouring for showing text diffs.
 package colour
 
+import (
+	"os"
+	"sync"
+)
+
 // ANSI codes for coloured output, they are all the same length so as not to throw off
 // alignment of [text/tabwriter].
 const (
@@ -10,6 +15,19 @@ const (
 	codeReset  = "\x1b[000000m" // Reset all attributes
 )
 
+// getColourOnce is a [sync.OnceValues] function that returns the state of
+// $NO_COLOR and $FORCE_COLOR, once and only once to avoid us calling
+// os.Getenv on every call to a colour function.
+var getColourOnce = sync.OnceValues(getColour)
+
+// getColour returns whether $NO_COLOR and $FORCE_COLOR were set.
+func getColour() (noColour bool, forceColour bool) {
+	no := os.Getenv("NO_COLOR") != ""
+	force := os.Getenv("FORCE_COLOR") != ""
+
+	return no, force
+}
+
 // Header returns a diff header styled string.
 func Header(text string) string {
 	return sprint(codeHeader, text)
@@ -29,5 +47,16 @@ func Red(text string) string {
 //
 // It handles checking for NO_COLOR and FORCE_COLOR.
 func sprint(code, text string) string {
+	noColor, forceColor := getColourOnce()
+
+	// $FORCE_COLOR overrides $NO_COLOR
+	if forceColor {
+		return code + text + codeReset
+	}
+
+	// $NO_COLOR is next
+	if noColor {
+		return text
+	}
 	return code + text + codeReset
 }