From b04103fa1d8ffacb8b1f8b2c99b52c72b1d24d5f Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Mon, 17 Mar 2025 11:06:34 +0100
Subject: [PATCH] [red-knot] Add `--color` CLI option (#16758)

## Summary

This PR adds a new `--color` CLI option that controls whether the output
should be colorized or not.

This is implements part of
https://github.com/astral-sh/ruff/issues/16727 except that it doesn't
implement the persistent configuration support as initially proposed in
the CLI document. I realized, that having this as a persistent
configuration is somewhat awkward because we may end up writing tracing
logs **before** we loaded and resolved the settings. Arguably, it's
probably fine to color the output up to that point, but it feels like a
somewhat broken experience. That's why I decided not to add the
persistent configuration option for now.


## Test Plan

I tested this change manually by running Red Knot with `--color=always`,
`--color=never`, and `--color=auto` (or no argument) and verified that:

* The diagnostics are or aren't colored
* The tracing output is or isn't colored.

---------

Co-authored-by: David Peter <sharkdp@users.noreply.github.com>
---
 crates/red_knot/src/args.rs | 18 ++++++++++++++++++
 crates/red_knot/src/main.rs | 36 ++++++++++++++++++++++++++++--------
 2 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/crates/red_knot/src/args.rs b/crates/red_knot/src/args.rs
index ffcf5ca911..8011bc5cd4 100644
--- a/crates/red_knot/src/args.rs
+++ b/crates/red_knot/src/args.rs
@@ -79,6 +79,10 @@ pub(crate) struct CheckCommand {
     #[arg(long)]
     pub(crate) output_format: Option<OutputFormat>,
 
+    /// Control when colored output is used.
+    #[arg(long, value_name = "WHEN")]
+    pub(crate) color: Option<TerminalColor>,
+
     /// Use exit code 1 if there are any warning-level diagnostics.
     #[arg(long, conflicts_with = "exit_zero", default_missing_value = "true", num_args=0..1)]
     pub(crate) error_on_warning: Option<bool>,
@@ -247,3 +251,17 @@ impl From<OutputFormat> for ruff_db::diagnostic::DiagnosticFormat {
         }
     }
 }
+
+/// Control when colored output is used.
+#[derive(Copy, Clone, Hash, Debug, PartialEq, Eq, PartialOrd, Ord, Default, clap::ValueEnum)]
+pub(crate) enum TerminalColor {
+    /// Display colors if the output goes to an interactive terminal.
+    #[default]
+    Auto,
+
+    /// Always display colors.
+    Always,
+
+    /// Never display colors.
+    Never,
+}
diff --git a/crates/red_knot/src/main.rs b/crates/red_knot/src/main.rs
index f6139358ad..9452fbf748 100644
--- a/crates/red_knot/src/main.rs
+++ b/crates/red_knot/src/main.rs
@@ -4,7 +4,7 @@ use std::process::{ExitCode, Termination};
 use anyhow::Result;
 use std::sync::Mutex;
 
-use crate::args::{Args, CheckCommand, Command};
+use crate::args::{Args, CheckCommand, Command, TerminalColor};
 use crate::logging::setup_tracing;
 use anyhow::{anyhow, Context};
 use clap::Parser;
@@ -76,6 +76,8 @@ pub(crate) fn version() -> Result<()> {
 }
 
 fn run_check(args: CheckCommand) -> anyhow::Result<ExitStatus> {
+    set_colored_override(args.color);
+
     let verbosity = args.verbosity.level();
     countme::enable(verbosity.is_trace());
     let _guard = setup_tracing(verbosity)?;
@@ -257,16 +259,16 @@ impl MainLoop {
                     result,
                     revision: check_revision,
                 } => {
+                    let terminal_settings = db.project().settings(db).terminal();
                     let display_config = DisplayDiagnosticConfig::default()
-                        .format(db.project().settings(db).terminal().output_format)
+                        .format(terminal_settings.output_format)
                         .color(colored::control::SHOULD_COLORIZE.should_colorize());
 
-                    let min_error_severity =
-                        if db.project().settings(db).terminal().error_on_warning {
-                            Severity::Warning
-                        } else {
-                            Severity::Error
-                        };
+                    let min_error_severity = if terminal_settings.error_on_warning {
+                        Severity::Warning
+                    } else {
+                        Severity::Error
+                    };
 
                     if check_revision == revision {
                         if db.project().files(db).is_empty() {
@@ -363,3 +365,21 @@ enum MainLoopMessage {
     ApplyChanges(Vec<watch::ChangeEvent>),
     Exit,
 }
+
+fn set_colored_override(color: Option<TerminalColor>) {
+    let Some(color) = color else {
+        return;
+    };
+
+    match color {
+        TerminalColor::Auto => {
+            colored::control::unset_override();
+        }
+        TerminalColor::Always => {
+            colored::control::set_override(true);
+        }
+        TerminalColor::Never => {
+            colored::control::set_override(false);
+        }
+    }
+}