diff --git a/doc/user-guide/src/environment-variables.md b/doc/user-guide/src/environment-variables.md index 8ef3351945..97b876c508 100644 --- a/doc/user-guide/src/environment-variables.md +++ b/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. diff --git a/src/currentprocess/filesource.rs b/src/currentprocess/filesource.rs index 36aa508b63..0f021fb212 100644 --- a/src/currentprocess/filesource.rs +++ b/src/currentprocess/filesource.rs @@ -198,7 +198,7 @@ impl Write for TestWriterLock<'_> { #[cfg(feature = "test")] pub(super) type TestWriterInner = Arc>>; /// 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); diff --git a/src/currentprocess/terminalsource.rs b/src/currentprocess/terminalsource.rs index 2a6e28b97b..19658ce8c5 100644 --- a/src/currentprocess/terminalsource.rs +++ b/src/currentprocess/terminalsource.rs @@ -11,7 +11,7 @@ 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 { @@ -19,6 +19,8 @@ pub(super) enum StreamSelector { 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,7 +169,7 @@ impl ColorableTerminal { s.set_color(spec) } #[cfg(feature = "test")] - TerminalInner::TestWriter(_) => Ok(()), + TerminalInner::TestWriter(_, _) => Ok(()), } } @@ -165,7 +177,7 @@ impl ColorableTerminal { match self.inner.lock().unwrap().deref_mut() { TerminalInner::StandardStream(s, _color) => s.reset(), #[cfg(feature = "test")] - TerminalInner::TestWriter(_) => Ok(()), + TerminalInner::TestWriter(_, _) => Ok(()), } } @@ -173,7 +185,7 @@ impl ColorableTerminal { 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,7 +202,7 @@ 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), } } @@ -198,7 +210,7 @@ impl io::Write for ColorableTerminal { 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, + ); + } +}