Add toolchain override support

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,8 @@ 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"),
         )
         .subcommand(
             SubCommand::with_name("dump-testament")
@@ -852,8 +854,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

@@ -40,6 +40,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 +105,7 @@ impl Cfg {
             temp_cfg,
             gpg_key,
             notify_handler,
+            toolchain_override: None,
             env_override,
             dist_root_url: dist_root,
             dist_root_server,
@@ -144,6 +146,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 +287,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::Environment));
+        }
+
+        // Check RUSTUP_TOOLCHAIN
         if let Some(ref name) = self.env_override {
             override_ = Some((name.to_string(), OverrideReason::Environment));
         }

tests/cli-misc.rs

@@ -905,3 +905,69 @@ fn which() {
         );
     });
 }
+
+#[test]
+fn override_by_toolchain_with_which() {
+    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_err(
+            config,
+            &["rustup", "+foo", "which", "rustc"],
+            "toolchain 'foo' is not installed",
+        );
+        expect_stderr_ok(
+            config,
+            &["rustup", "+stable", "set", "profile", "minimal"],
+            "profile set to 'minimal'",
+        );
+        expect_stdout_ok(config, &["rustup", "default"], "nightly-x86_64-");
+    });
+}