Add toolchain override support

Author: BeniCheni

src/cli/rustup_mode.rs

@@ -34,6 +34,10 @@ pub fn main() -> Result<()> {
     let quiet = matches.is_present("quiet");
     let cfg = &common::set_globals(verbose, quiet)?;
 
+    if matches.is_present("+toolchain") {
+        override_by_toolchain(cfg, &matches)?;
+    }
+
     if maybe_upgrade_data(cfg, &matches)? {
         return Ok(());
     }
@@ -130,10 +134,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")
@@ -819,8 +821,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")
@@ -1148,6 +1149,29 @@ fn toolchain_remove(cfg: &Cfg, m: &ArgMatches<'_>) -> Result<()> {
     Ok(())
 }
 
+fn override_by_toolchain(cfg: &Cfg, m: &ArgMatches<'_>) -> Result<()> {
+    let toolchain = m.value_of("+toolchain").expect("");
+    let toolchain = cfg.get_toolchain(&toolchain[1..], false)?;
+
+    let status = if !toolchain.is_custom() {
+        Some(toolchain.install_from_dist_if_not_installed()?)
+    } else if !toolchain.exists() {
+        return Err(ErrorKind::ToolchainNotInstalled(toolchain.name().to_string()).into());
+    } else {
+        None
+    };
+
+    let path = utils::current_dir()?;
+    toolchain.make_override(&path)?;
+
+    if let Some(status) = status {
+        println!();
+        common::show_channel_update(cfg, toolchain.name(), Ok(status))?;
+    }
+
+    Ok(())
+}
+
 fn override_add(cfg: &Cfg, m: &ArgMatches<'_>) -> Result<()> {
     let toolchain = m.value_of("toolchain").expect("");
     let toolchain = cfg.get_toolchain(toolchain, false)?;

tests/cli-misc.rs

@@ -905,3 +905,68 @@ 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_stderr_ok(
+            config,
+            &["rustup", "+stable", "which", "rustc"],
+            "set to 'stable-x86_64-",
+        );
+        expect_err(
+            config,
+            &["rustup", "+foo", "which", "rustc"],
+            "toolchain 'foo' is not installed",
+        );
+    });
+}