Add --color cli option to force colored output

crates/ruff/src/args.rs

@@ -67,6 +67,15 @@ pub struct GlobalConfigArgs {
     // we emit an error later on, after the initial parsing by clap.
     #[arg(long, help_heading = "Global options", global = true)]
     pub isolated: bool,
+
+    /// Control when colored output is used.
+    #[arg(
+        long,
+        value_name = "WHEN",
+        help_heading = "Global options",
+        global = true
+    )]
+    pub(crate) color: Option<TerminalColor>,
 }
 
 impl GlobalConfigArgs {
@@ -80,6 +89,20 @@ impl GlobalConfigArgs {
     }
 }
 
+/// 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,
+}
+
 // Configures Clap v3-style help menu colors
 const STYLES: Styles = Styles::styled()
     .header(AnsiColor::Green.on_default().effects(Effects::BOLD))

crates/ruff/src/lib.rs

@@ -1,5 +1,6 @@
 #![allow(clippy::print_stdout)]
 
+use std::ffi::OsString;
 use std::fs::File;
 use std::io::{self, BufWriter, Write, stdout};
 use std::path::{Path, PathBuf};
@@ -21,7 +22,7 @@ use ruff_linter::{fs, warn_user, warn_user_once};
 use ruff_workspace::Settings;
 
 use crate::args::{
-    AnalyzeCommand, AnalyzeGraphCommand, Args, CheckCommand, Command, FormatCommand,
+    AnalyzeCommand, AnalyzeGraphCommand, Args, CheckCommand, Command, FormatCommand, TerminalColor,
 };
 use crate::printer::{Flags as PrinterFlags, Printer};
 
@@ -131,6 +132,13 @@ pub fn run(
         global_options,
     }: Args,
 ) -> Result<ExitStatus> {
+    // Set color before so all outputs are properly colored
+    if let Some(color_override) =
+        colored_override(global_options.color, std::env::var_os("FORCE_COLOR"))
+    {
+        colored::control::set_override(color_override);
+    }
+
     {
         ruff_db::set_program_version(crate::version::version().to_string()).unwrap();
         let default_panic_hook = std::panic::take_hook();
@@ -514,6 +522,22 @@ https://github.com/astral-sh/ruff/issues/new?title=%5BLinter%20panic%5D
     Ok(ExitStatus::Success)
 }
 
+fn colored_override(
+    color: Option<TerminalColor>,
+    env_force_color: Option<OsString>,
+) -> Option<bool> {
+    match color {
+        // Cli arguments should take precedence over env vars.
+        Some(TerminalColor::Always) => Some(true),
+        Some(TerminalColor::Never) => Some(false),
+        // Default to no override, but respect FORCE_COLOR.
+        Some(TerminalColor::Auto) | None => {
+            // support FORCE_COLOR env var
+            env_force_color.map(|force_color: OsString| !force_color.is_empty())
+        }
+    }
+}
+
 #[cfg(test)]
 mod test_file_change_detector {
     use std::path::PathBuf;
@@ -634,3 +658,26 @@ mod test_file_change_detector {
         );
     }
 }
+
+#[cfg(test)]
+mod test_set_colored_override {
+    use crate::{args::TerminalColor, colored_override};
+
+    #[test]
+    fn force_color_env_is_respected() {
+        assert_eq!(colored_override(None, Some("1".into())), Some(true));
+    }
+
+    #[test]
+    fn cli_args_takes_precedences_over_force_color_env() {
+        assert_eq!(
+            colored_override(Some(TerminalColor::Never), Some("1".into())),
+            Some(false)
+        );
+
+        assert_eq!(
+            colored_override(Some(TerminalColor::Always), None),
+            Some(true)
+        );
+    }
+}

crates/ruff/src/main.rs

@@ -32,13 +32,6 @@ pub fn main() -> ExitCode {
     #[cfg(windows)]
     assert!(colored::control::set_virtual_terminal(true).is_ok());
 
-    // support FORCE_COLOR env var
-    if let Some(force_color) = std::env::var_os("FORCE_COLOR") {
-        if !force_color.is_empty() {
-            colored::control::set_override(true);
-        }
-    }
-
     let args = wild::args_os();
     let args = match argfile::expand_args_from(args, argfile::parse_fromfile, argfile::PREFIX)
         .context("Failed to read CLI arguments from files")

docs/configuration.md

@@ -540,7 +540,7 @@ Commands:
   help     Print this message or the help of the given subcommand(s)
 
 Options:
-  -h, --help     Print help
+  -h, --help     Print help (see more with '--help')
   -V, --version  Print version
 
 Log levels:
@@ -559,6 +559,9 @@ Global options:
           configuration files that were also specified using `--config`
       --isolated
           Ignore all configuration files
+      --color <WHEN>
+          Control when colored output is used [possible values: auto, always,
+          never]
 
 For help with a specific command, see: `ruff help <command>`.
 ```
@@ -628,7 +631,7 @@ Options:
       --show-settings
           See the settings Ruff will use to lint a given Python file
   -h, --help
-          Print help
+          Print help (see more with '--help')
 
 Rule selection:
       --select <RULE_CODE>
@@ -696,6 +699,9 @@ Global options:
           configuration files that were also specified using `--config`
       --isolated
           Ignore all configuration files
+      --color <WHEN>
+          Control when colored output is used [possible values: auto, always,
+          never]
 ```
 
 <!-- End auto-generated check help. -->
@@ -786,6 +792,9 @@ Global options:
           configuration files that were also specified using `--config`
       --isolated
           Ignore all configuration files
+      --color <WHEN>
+          Control when colored output is used [possible values: auto, always,
+          never]
 ```
 
 <!-- End auto-generated format help. -->