Merge pull request #3435 from rami3l/feat/force-color-output

Support `RUSTUP_TERM_COLOR` as an override environment variable

Author: rbtcollins

doc/user-guide/src/environment-variables.md

@@ -29,6 +29,10 @@
   determines the directory that traces will be written too. Traces are of the
   form PID.trace. Traces can be read by the Catapult project [tracing viewer].
 
+- `RUSTUP_TERM_COLOR` (default: `auto`) Controls whether colored output is used in the terminal.
+  Set to `auto` to use colors only in tty streams, to `always` to always enable colors,
+  or to `never` to disable colors.
+
 - `RUSTUP_UNPACK_RAM` *unstable* (default free memory or 500MiB if unable to tell, min 210MiB) Caps the amount of
   RAM `rustup` will use for IO tasks while unpacking.
 

src/currentprocess/filesource.rs

@@ -198,7 +198,7 @@ impl Write for TestWriterLock<'_> {
 #[cfg(feature = "test")]
 pub(super) type TestWriterInner = Arc<Mutex<Vec<u8>>>;
 /// A thread-safe test file handle that pretends to be e.g. stdout.
-#[derive(Clone)]
+#[derive(Clone, Default)]
 #[cfg(feature = "test")]
 pub(super) struct TestWriter(TestWriterInner);
 

src/currentprocess/terminalsource.rs

@@ -11,14 +11,16 @@ use termcolor::{ColorChoice, ColorSpec, StandardStream, StandardStreamLock, Writ
 
 #[cfg(feature = "test")]
 use super::filesource::{TestWriter, TestWriterLock};
-use super::process;
+use super::{process, varsource::VarSource};
 
 /// Select what stream to make a terminal on
 pub(super) enum StreamSelector {
     Stdout,
     Stderr,
     #[cfg(feature = "test")]
     TestWriter(TestWriter),
+    #[cfg(feature = "test")]
+    TestTtyWriter(TestWriter),
 }
 
 impl StreamSelector {
@@ -36,6 +38,8 @@ impl StreamSelector {
             },
             #[cfg(feature = "test")]
             StreamSelector::TestWriter(_) => false,
+            #[cfg(feature = "test")]
+            StreamSelector::TestTtyWriter(_) => true,
         }
     }
 }
@@ -55,7 +59,7 @@ pub struct ColorableTerminal {
 enum TerminalInner {
     StandardStream(StandardStream, ColorSpec),
     #[cfg(feature = "test")]
-    TestWriter(TestWriter),
+    TestWriter(TestWriter, ColorChoice),
 }
 
 pub struct ColorableTerminalLocked {
@@ -73,14 +77,20 @@ enum TerminalInnerLocked {
 }
 
 impl ColorableTerminal {
-    /// Construct a terminal for the selected stream. Colors written to the
-    /// terminal will be discarded if the stream is not a tty.
+    /// A terminal that supports colorisation of a stream.
+    /// If `RUSTUP_TERM_COLOR` is set to `always`, or if the stream is a tty and
+    /// `RUSTUP_TERM_COLOR` either unset or set to `auto`,
+    /// then color commands will be sent to the stream.
+    /// Otherwise color commands are discarded.
     pub(super) fn new(stream: StreamSelector) -> Self {
-        let is_a_tty = stream.is_a_tty();
-        let choice = if is_a_tty {
-            ColorChoice::Auto
-        } else {
-            ColorChoice::Never
+        let env_override = process()
+            .var("RUSTUP_TERM_COLOR")
+            .map(|it| it.to_lowercase());
+        let choice = match env_override.as_deref() {
+            Ok("always") => ColorChoice::Always,
+            Ok("never") => ColorChoice::Never,
+            _ if stream.is_a_tty() => ColorChoice::Auto,
+            _ => ColorChoice::Never,
         };
         let inner = match stream {
             StreamSelector::Stdout => {
@@ -90,7 +100,9 @@ impl ColorableTerminal {
                 TerminalInner::StandardStream(StandardStream::stderr(choice), ColorSpec::new())
             }
             #[cfg(feature = "test")]
-            StreamSelector::TestWriter(w) => TerminalInner::TestWriter(w),
+            StreamSelector::TestWriter(w) | StreamSelector::TestTtyWriter(w) => {
+                TerminalInner::TestWriter(w, choice)
+            }
         };
         ColorableTerminal {
             inner: Arc::new(Mutex::new(inner)),
@@ -118,7 +130,7 @@ impl ColorableTerminal {
                     TerminalInnerLocked::StandardStream(locked)
                 }
                 #[cfg(feature = "test")]
-                TerminalInner::TestWriter(w) => TerminalInnerLocked::TestWriter(w.lock()),
+                TerminalInner::TestWriter(w, _) => TerminalInnerLocked::TestWriter(w.lock()),
             });
             // ColorableTerminalLocked { inner, guard, locked }
             uninit.assume_init()
@@ -132,7 +144,7 @@ impl ColorableTerminal {
                 s.set_color(spec)
             }
             #[cfg(feature = "test")]
-            TerminalInner::TestWriter(_) => Ok(()),
+            TerminalInner::TestWriter(_, _) => Ok(()),
         }
     }
 
@@ -143,7 +155,7 @@ impl ColorableTerminal {
                 s.set_color(spec)
             }
             #[cfg(feature = "test")]
-            TerminalInner::TestWriter(_) => Ok(()),
+            TerminalInner::TestWriter(_, _) => Ok(()),
         }
     }
 
@@ -157,23 +169,23 @@ impl ColorableTerminal {
                 s.set_color(spec)
             }
             #[cfg(feature = "test")]
-            TerminalInner::TestWriter(_) => Ok(()),
+            TerminalInner::TestWriter(_, _) => Ok(()),
         }
     }
 
     pub fn reset(&mut self) -> io::Result<()> {
         match self.inner.lock().unwrap().deref_mut() {
             TerminalInner::StandardStream(s, _color) => s.reset(),
             #[cfg(feature = "test")]
-            TerminalInner::TestWriter(_) => Ok(()),
+            TerminalInner::TestWriter(_, _) => Ok(()),
         }
     }
 
     pub fn carriage_return(&mut self) -> io::Result<()> {
         match self.inner.lock().unwrap().deref_mut() {
             TerminalInner::StandardStream(s, _color) => s.write(b"\r")?,
             #[cfg(feature = "test")]
-            TerminalInner::TestWriter(w) => w.write(b"\r")?,
+            TerminalInner::TestWriter(w, _) => w.write(b"\r")?,
         };
         Ok(())
     }
@@ -190,15 +202,15 @@ impl io::Write for ColorableTerminal {
         match self.inner.lock().unwrap().deref_mut() {
             TerminalInner::StandardStream(s, _) => s.write(buf),
             #[cfg(feature = "test")]
-            TerminalInner::TestWriter(w) => w.write(buf),
+            TerminalInner::TestWriter(w, _) => w.write(buf),
         }
     }
 
     fn flush(&mut self) -> std::result::Result<(), io::Error> {
         match self.inner.lock().unwrap().deref_mut() {
             TerminalInner::StandardStream(s, _) => s.flush(),
             #[cfg(feature = "test")]
-            TerminalInner::TestWriter(w) => w.flush(),
+            TerminalInner::TestWriter(w, _) => w.flush(),
         }
     }
 }
@@ -220,3 +232,56 @@ impl io::Write for ColorableTerminalLocked {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use rustup_macros::unit_test as test;
+
+    use super::*;
+    use crate::{currentprocess, test::Env};
+
+    #[test]
+    fn term_color_choice() {
+        fn assert_color_choice(env_val: &str, stream: StreamSelector, color_choice: ColorChoice) {
+            let mut vars = HashMap::new();
+            vars.env("RUSTUP_TERM_COLOR", env_val);
+            let tp = currentprocess::TestProcess {
+                vars,
+                ..Default::default()
+            };
+            currentprocess::with(tp.into(), || {
+                let term = ColorableTerminal::new(stream);
+                let inner = term.inner.lock().unwrap();
+                assert!(matches!(
+                    &*inner,
+                    &TerminalInner::TestWriter(_, choice) if choice == color_choice
+                ));
+            });
+        }
+
+        assert_color_choice(
+            "aLWayS",
+            StreamSelector::TestWriter(Default::default()),
+            ColorChoice::Always,
+        );
+        assert_color_choice(
+            "neVer",
+            StreamSelector::TestWriter(Default::default()),
+            ColorChoice::Never,
+        );
+        // tty + `auto` enables the colors.
+        assert_color_choice(
+            "AutO",
+            StreamSelector::TestTtyWriter(Default::default()),
+            ColorChoice::Auto,
+        );
+        // non-tty + `auto` does not enable the colors.
+        assert_color_choice(
+            "aUTo",
+            StreamSelector::TestWriter(Default::default()),
+            ColorChoice::Never,
+        );
+    }
+}