Add auto color detection

README.md

@@ -145,6 +145,26 @@ let output = try Chroma.highlight(
 )
 ```
 
+### Color Output
+
+Chroma defaults to auto-detecting ANSI output based on TTY, `TERM=dumb`, and common environment variables (`NO_COLOR`, `CHROMA_NO_COLOR`, `FORCE_COLOR`).
+
+```swift
+// Force ANSI output
+let output = try Chroma.highlight(
+    code,
+    language: .swift,
+    options: .init(colorMode: .always)
+)
+
+// Auto-detect for stderr output
+let output = try Chroma.highlight(
+    code,
+    language: .swift,
+    options: .init(colorMode: .auto(output: .stderr))
+)
+```
+
 ### Output Indentation
 
 Indent the entire output by a specified number of spaces:

Sources/Chroma/ColorMode.swift

@@ -0,0 +1,101 @@
+import Foundation
+import Rainbow
+#if os(Linux)
+import Glibc
+#else
+import Darwin
+#endif
+
+/// Output stream used for auto color detection.
+public enum ColorOutput: Equatable {
+    case stdout
+    case stderr
+}
+
+/// Controls when ANSI colors are emitted.
+public enum ColorMode: Equatable {
+    /// Decide automatically using TTY detection, TERM, and color-related environment variables.
+    case auto(output: ColorOutput)
+    /// Always emit ANSI colors.
+    case always
+    /// Never emit ANSI colors.
+    case never
+    /// Defer to `Rainbow.enabled`.
+    case inheritRainbowEnabled
+}
+
+extension ColorMode {
+    /// Resolve whether ANSI colors should be enabled in the current process environment.
+    public func isEnabled() -> Bool {
+        resolve(
+            rainbowEnabled: Rainbow.enabled,
+            environment: ProcessInfo.processInfo.environment,
+            isTTY: Self.isTTY
+        )
+    }
+
+    func resolve(
+        rainbowEnabled: Bool,
+        environment: [String: String],
+        isTTY: (ColorOutput) -> Bool
+    ) -> Bool {
+        switch self {
+        case .always:
+            return true
+        case .never:
+            return false
+        case .inheritRainbowEnabled:
+            return rainbowEnabled
+        case let .auto(output):
+            return ColorEnvironment.shouldEnableColor(
+                output: output,
+                rainbowEnabled: rainbowEnabled,
+                environment: environment,
+                isTTY: isTTY
+            )
+        }
+    }
+
+    private static func isTTY(_ output: ColorOutput) -> Bool {
+        switch output {
+        case .stdout:
+            return isatty(fileno(stdout)) != 0
+        case .stderr:
+            return isatty(fileno(stderr)) != 0
+        }
+    }
+}
+
+private enum ColorEnvironment {
+    static func shouldEnableColor(
+        output: ColorOutput,
+        rainbowEnabled: Bool,
+        environment: [String: String],
+        isTTY: (ColorOutput) -> Bool
+    ) -> Bool {
+        guard rainbowEnabled else { return false }
+        if hasNonEmptyValue("CHROMA_NO_COLOR", environment: environment) {
+            return false
+        }
+        if hasTruthyValue("FORCE_COLOR", environment: environment) {
+            return true
+        }
+        if hasNonEmptyValue("NO_COLOR", environment: environment) {
+            return false
+        }
+        if let term = environment["TERM"], term.lowercased() == "dumb" {
+            return false
+        }
+        return isTTY(output)
+    }
+
+    private static func hasNonEmptyValue(_ key: String, environment: [String: String]) -> Bool {
+        guard let value = environment[key] else { return false }
+        return !value.isEmpty
+    }
+
+    private static func hasTruthyValue(_ key: String, environment: [String: String]) -> Bool {
+        guard let value = environment[key] else { return false }
+        return !value.isEmpty && value != "0"
+    }
+}

Sources/Chroma/HighlightOptions.swift

@@ -49,6 +49,8 @@ public struct HighlightOptions: Equatable {
     }
 
     public var theme: Theme?
+    /// Controls whether ANSI colors are emitted.
+    public var colorMode: ColorMode
     public var missingLanguageHandling: MissingLanguageHandling
     public var diff: DiffHighlight
     public var highlightLines: LineRangeSet
@@ -63,13 +65,15 @@ public struct HighlightOptions: Equatable {
 
     public init(
         theme: Theme? = nil,
+        colorMode: ColorMode = .auto(output: .stdout),
         missingLanguageHandling: MissingLanguageHandling = .error,
         diff: DiffHighlight = .auto(),
         highlightLines: LineRangeSet = .init(),
         lineNumbers: LineNumberOptions = .none,
         indent: Int = 0
     ) {
         self.theme = theme
+        self.colorMode = colorMode
         self.missingLanguageHandling = missingLanguageHandling
         self.diff = diff
         self.highlightLines = highlightLines

Sources/Chroma/Renderer.swift

@@ -29,7 +29,8 @@ final class Renderer {
         estimatedSegments: Int?,
         tokenStream: (_ emit: (Token) -> Void) -> Void
     ) -> String {
-        if !Rainbow.enabled &&
+        let colorEnabled = options.colorMode.isEnabled()
+        if !colorEnabled &&
             options.highlightLines.ranges.isEmpty &&
             options.indent == 0 &&
             !options.lineNumbers.isEnabled {
@@ -46,7 +47,7 @@ final class Renderer {
         var writer = AnsiWriter(
             estimatedTextLength: code.count,
             estimatedSegments: estimatedSegments ?? max(16, code.count / 8),
-            isEnabled: Rainbow.enabled
+            isEnabled: colorEnabled
         )
 
         var atLineStart = true

Sources/ChromaDemo/ChromaDemo.swift

@@ -1562,8 +1562,7 @@ struct DemoArguments {
     }
 
     var theme: ThemeChoice = .dark
-    var forceColor: Bool = false
-    var noColor: Bool = false
+    var colorMode: ColorMode = .auto(output: .stdout)
     var listLanguages: Bool = false
     var lang: String?
     var showComponents: Bool = true
@@ -1590,9 +1589,9 @@ struct DemoArguments {
             case "--light":
                 parsed.theme = .light
             case "--force-color":
-                parsed.forceColor = true
+                parsed.colorMode = .always
             case "--no-color":
-                parsed.noColor = true
+                parsed.colorMode = .never
             case "--list-languages":
                 parsed.listLanguages = true
             case "--lang":
@@ -1638,11 +1637,8 @@ struct ChromaDemo {
             let args = Array(CommandLine.arguments.dropFirst())
             let options = try DemoArguments.parse(args)
 
-            if options.noColor {
-                Rainbow.enabled = false
-            } else if options.forceColor {
-                Rainbow.enabled = true
-            }
+            let colorMode = options.colorMode
+            Rainbow.enabled = colorMode.isEnabled()
 
             let theme: Theme = (options.theme == .light) ? .light : .dark
             let registry = LanguageRegistry.builtIn()
@@ -1663,6 +1659,7 @@ struct ChromaDemo {
                 registry: registry,
                 highlighter: highlighter,
                 theme: theme,
+                colorMode: colorMode,
                 showHeader: options.showComponents
             )
         } catch DemoError.invalidArguments(let message) {
@@ -1696,6 +1693,7 @@ struct ChromaDemo {
         registry: LanguageRegistry,
         highlighter: Highlighter,
         theme: Theme,
+        colorMode: ColorMode,
         showHeader: Bool
     ) throws {
         let displayName = registry.language(for: demo.id)?.displayName ?? demo.id.rawValue
@@ -1707,14 +1705,19 @@ struct ChromaDemo {
         }
 
         printSection("Sample")
-        print(try highlighter.highlight(demo.sample, language: demo.id))
+        print(try highlighter.highlight(
+            demo.sample,
+            language: demo.id,
+            options: .init(colorMode: colorMode)
+        ))
         print("")
 
         printSection("Sample + Line Numbers + Highlights")
         print(try highlighter.highlight(
             demo.sample,
             language: demo.id,
             options: .init(
+                colorMode: colorMode,
                 highlightLines: demo.highlightLines,
                 lineNumbers: .init(start: 1)
             )
@@ -1726,6 +1729,7 @@ struct ChromaDemo {
             demo.patch,
             language: demo.id,
             options: .init(
+                colorMode: colorMode,
                 diff: .patch(style: .background(diffCode: .plain), presentation: .verbose),
                 lineNumbers: .init(start: 1)
             )
@@ -1737,6 +1741,7 @@ struct ChromaDemo {
             demo.patch,
             language: demo.id,
             options: .init(
+                colorMode: colorMode,
                 diff: .patch(style: .foreground(contextCode: .plain), presentation: .compact),
                 lineNumbers: .init(start: 1)
             )

Tests/ChromaTests/ChromaFacadeTests.swift

@@ -10,7 +10,7 @@ struct ChromaFacadeTests {
         let output = try Chroma.highlight(
             "let value = 1",
             language: .swift,
-            options: .init(theme: theme)
+            options: .init(theme: theme, colorMode: .always)
         )
 
         let expected = renderExpected([

Tests/ChromaTests/ColorModeTests.swift

@@ -0,0 +1,75 @@
+import Testing
+@testable import Chroma
+
+@Suite("ColorMode")
+struct ColorModeTests {
+    @Test("Auto disables when NO_COLOR is set")
+    func autoDisablesOnNoColor() {
+        let enabled = ColorMode.auto(output: .stdout).resolve(
+            rainbowEnabled: true,
+            environment: ["NO_COLOR": "1"],
+            isTTY: { _ in true }
+        )
+        #expect(!enabled)
+    }
+
+    @Test("Auto disables when CHROMA_NO_COLOR is set")
+    func autoDisablesOnChromaNoColor() {
+        let enabled = ColorMode.auto(output: .stderr).resolve(
+            rainbowEnabled: true,
+            environment: ["CHROMA_NO_COLOR": "true"],
+            isTTY: { _ in true }
+        )
+        #expect(!enabled)
+    }
+
+    @Test("Auto disables for dumb terminals")
+    func autoDisablesOnDumbTerm() {
+        let enabled = ColorMode.auto(output: .stdout).resolve(
+            rainbowEnabled: true,
+            environment: ["TERM": "dumb"],
+            isTTY: { _ in true }
+        )
+        #expect(!enabled)
+    }
+
+    @Test("Auto disables when not a TTY")
+    func autoDisablesOnNonTTY() {
+        let enabled = ColorMode.auto(output: .stdout).resolve(
+            rainbowEnabled: true,
+            environment: [:],
+            isTTY: { _ in false }
+        )
+        #expect(!enabled)
+    }
+
+    @Test("Auto respects FORCE_COLOR even without TTY")
+    func autoHonorsForceColor() {
+        let enabled = ColorMode.auto(output: .stdout).resolve(
+            rainbowEnabled: true,
+            environment: ["FORCE_COLOR": "1"],
+            isTTY: { _ in false }
+        )
+        #expect(enabled)
+    }
+
+    @Test("CHROMA_NO_COLOR overrides FORCE_COLOR")
+    func chromaNoColorOverridesForceColor() {
+        let enabled = ColorMode.auto(output: .stdout).resolve(
+            rainbowEnabled: true,
+            environment: ["CHROMA_NO_COLOR": "1", "FORCE_COLOR": "1"],
+            isTTY: { _ in true }
+        )
+        #expect(!enabled)
+    }
+
+    @Test("Auto respects Rainbow.enabled")
+    func autoRespectsRainbowEnabled() {
+        let enabled = ColorMode.auto(output: .stdout).resolve(
+            rainbowEnabled: false,
+            environment: [:],
+            isTTY: { _ in true }
+        )
+        #expect(!enabled)
+    }
+}

Tests/ChromaTests/HighlightOptionsTests.swift

@@ -7,6 +7,7 @@ struct HighlightOptionsTests {
     func defaults() {
         let options = HighlightOptions()
         #expect(options.theme == nil)
+        #expect(options.colorMode == .auto(output: .stdout))
         #expect(options.missingLanguageHandling == .error)
         #expect(options.diff == .auto())
         #expect(options.highlightLines == LineRangeSet())