Add toolchain override support on command line

Author: BeniCheni

src/cli/rustup_mode.rs

@@ -35,6 +35,10 @@ pub fn main() -> Result<()> {
     let quiet = matches.is_present("quiet");
     let cfg = &mut common::set_globals(verbose, quiet)?;
 
+    if let Some(t) = matches.value_of("+toolchain") {
+        cfg.set_toolchain_override(&t[1..]);
+    }
+
     if maybe_upgrade_data(cfg, &matches)? {
         return Ok(());
     }
@@ -131,10 +135,15 @@ pub fn cli() -> App<'static, 'static> {
                 .long("quiet"),
         )
         .arg(
-            Arg::with_name("toolchain")
-                .help(TOOLCHAIN_ARG_HELP)
-                .long("toolchain")
-                .takes_value(true),
+            Arg::with_name("+toolchain")
+                .help("release channel (e.g. +stable) or custom toolchain to set override")
+                .validator(|s| {
+                    if s.starts_with('+') {
+                        Ok(())
+                    } else {
+                        Err("Toolchain overrides must begin with '+'".into())
+                    }
+                }),
         )
         .subcommand(
             SubCommand::with_name("dump-testament")
@@ -852,8 +861,7 @@ fn run(cfg: &Cfg, m: &ArgMatches<'_>) -> Result<()> {
 
 fn which(cfg: &Cfg, m: &ArgMatches<'_>) -> Result<()> {
     let binary = m.value_of("command").expect("");
-    let toolchain_provided = m.is_present("toolchain");
-    let binary_path = if toolchain_provided {
+    let binary_path = if m.is_present("toolchain") {
         let toolchain = m.value_of("toolchain").expect("");
         cfg.which_binary_by_toolchain(toolchain, binary)?
             .expect("binary not found")

src/config.rs

@@ -17,6 +17,7 @@ use crate::utils::utils;
 #[derive(Debug)]
 pub enum OverrideReason {
     Environment,
+    CommandLine,
     OverrideDB(PathBuf),
     ToolchainFile(PathBuf),
 }
@@ -25,6 +26,7 @@ impl Display for OverrideReason {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::result::Result<(), fmt::Error> {
         match self {
             Self::Environment => write!(f, "environment override by RUSTUP_TOOLCHAIN"),
+            Self::CommandLine => write!(f, "overridden by +toolchain on the command line"),
             Self::OverrideDB(path) => write!(f, "directory override for '{}'", path.display()),
             Self::ToolchainFile(path) => write!(f, "overridden by '{}'", path.display()),
         }
@@ -40,6 +42,7 @@ pub struct Cfg {
     pub download_dir: PathBuf,
     pub temp_cfg: temp::Cfg,
     pub gpg_key: Cow<'static, str>,
+    pub toolchain_override: Option<String>,
     pub env_override: Option<String>,
     pub dist_root_url: String,
     pub dist_root_server: String,
@@ -104,6 +107,7 @@ impl Cfg {
             temp_cfg,
             gpg_key,
             notify_handler,
+            toolchain_override: None,
             env_override,
             dist_root_url: dist_root,
             dist_root_server,
@@ -144,6 +148,10 @@ impl Cfg {
         Ok(())
     }
 
+    pub fn set_toolchain_override(&mut self, toolchain_override: &str) {
+        self.toolchain_override = Some(toolchain_override.to_owned());
+    }
+
     // Returns a profile, if one exists in the settings file.
     //
     // Returns `Err` if the settings file could not be read or the profile is
@@ -281,7 +289,12 @@ impl Cfg {
     pub fn find_override(&self, path: &Path) -> Result<Option<(Toolchain<'_>, OverrideReason)>> {
         let mut override_ = None;
 
-        // First check RUSTUP_TOOLCHAIN
+        // First check toolchain override from command
+        if let Some(ref name) = self.toolchain_override {
+            override_ = Some((name.to_string(), OverrideReason::CommandLine));
+        }
+
+        // Check RUSTUP_TOOLCHAIN
         if let Some(ref name) = self.env_override {
             override_ = Some((name.to_string(), OverrideReason::Environment));
         }
@@ -306,6 +319,10 @@ impl Cfg {
                     "the RUSTUP_TOOLCHAIN environment variable specifies an uninstalled toolchain"
                         .to_string()
                 }
+                OverrideReason::CommandLine => {
+                    "the +toolchain on the command line specifies an uninstalled toolchain"
+                        .to_string()
+                }
                 OverrideReason::OverrideDB(ref path) => format!(
                     "the directory override for '{}' specifies an uninstalled toolchain",
                     path.display()

tests/cli-misc.rs

@@ -905,3 +905,79 @@ fn which() {
         );
     });
 }
+
+#[test]
+fn override_by_toolchain_on_the_command_line() {
+    setup(&|config| {
+        #[cfg(windows)]
+        expect_stdout_ok(
+            config,
+            &["rustup", "+stable", "which", "rustc"],
+            "\\toolchains\\stable-x86_64-",
+        );
+        #[cfg(windows)]
+        expect_stdout_ok(
+            config,
+            &["rustup", "+stable", "which", "rustc"],
+            "\\bin\\rustc",
+        );
+        #[cfg(not(windows))]
+        expect_stdout_ok(
+            config,
+            &["rustup", "+stable", "which", "rustc"],
+            "/toolchains/stable-x86_64-",
+        );
+        #[cfg(not(windows))]
+        expect_stdout_ok(
+            config,
+            &["rustup", "+stable", "which", "rustc"],
+            "/bin/rustc",
+        );
+        expect_ok(config, &["rustup", "default", "nightly"]);
+        #[cfg(windows)]
+        expect_stdout_ok(
+            config,
+            &["rustup", "+nightly", "which", "rustc"],
+            "\\toolchains\\nightly-x86_64-",
+        );
+        #[cfg(windows)]
+        expect_stdout_ok(
+            config,
+            &["rustup", "+nightly", "which", "rustc"],
+            "\\bin\\rustc",
+        );
+        #[cfg(not(windows))]
+        expect_stdout_ok(
+            config,
+            &["rustup", "+nightly", "which", "rustc"],
+            "/toolchains/nightly-x86_64-",
+        );
+        #[cfg(not(windows))]
+        expect_stdout_ok(
+            config,
+            &["rustup", "+nightly", "which", "rustc"],
+            "/bin/rustc",
+        );
+        expect_stdout_ok(
+            config,
+            &["rustup", "+nightly", "show"],
+            "(overridden by +toolchain on the command line)",
+        );
+        expect_err(
+            config,
+            &["rustup", "+foo", "which", "rustc"],
+            "toolchain 'foo' is not installed",
+        );
+        expect_err(
+            config,
+            &["rustup", "@stable", "which", "rustc"],
+            "Invalid value for '<+toolchain>': Toolchain overrides must begin with '+'",
+        );
+        expect_stderr_ok(
+            config,
+            &["rustup", "+stable", "set", "profile", "minimal"],
+            "profile set to 'minimal'",
+        );
+        expect_stdout_ok(config, &["rustup", "default"], "nightly-x86_64-");
+    });
+}