Added support for bootstrapping with specific toolchain version (astr…

…al-sh#606)

CHANGELOG.md

@@ -20,6 +20,8 @@ _Unreleased_
 - Upgraded internal unearth dependency which resolved an issue where
   `rye add tensorflow` would not work.  #614
 
+- The installer now supports `RYE_TOOLCHAIN_VERSION`.  #606
+
 <!-- released start -->
 
 ## 0.21.0

docs/.includes/installer-options.md

@@ -5,3 +5,12 @@
     provided a suitable interpreter is automatically downloaded.
 
     At present only CPython 3.9 to 3.12 are supported.
+
+`RYE_TOOLCHAIN_VERSION`
+
+:   For Rye 0.22 and later a specific Python version can be picked rather
+    than the default.  This affects the internal toolchain version only.
+    It's useful for Docker builds where you can set the internal toolchain
+    to the same as your project to only fetch a single Python.
+
+    At present only CPython 3.9 to 3.12 are supported.

rye/src/bootstrap.rs

@@ -6,7 +6,7 @@ use std::process::Command;
 use std::sync::atomic::{self, AtomicBool};
 use std::{env, fs};
 
-use anyhow::{bail, Context, Error};
+use anyhow::{anyhow, bail, Context, Error};
 use console::style;
 use indicatif::{ProgressBar, ProgressStyle};
 use once_cell::sync::Lazy;
@@ -19,6 +19,7 @@ use crate::platform::{
     get_app_dir, get_canonical_py_path, get_toolchain_python_bin, list_known_toolchains,
     symlinks_supported,
 };
+use crate::pyproject::latest_available_python_version;
 use crate::sources::{get_download_url, PythonVersion, PythonVersionRequest};
 use crate::utils::{
     check_checksum, get_venv_python_bin, set_proxy_variables, symlink_file, unpack_archive,
@@ -71,6 +72,14 @@ fn is_up_to_date() -> bool {
 
 /// Bootstraps the venv for rye itself
 pub fn ensure_self_venv(output: CommandOutput) -> Result<PathBuf, Error> {
+    ensure_self_venv_with_toolchain(output, None)
+}
+
+/// Bootstraps the venv for rye itself
+pub fn ensure_self_venv_with_toolchain(
+    output: CommandOutput,
+    toolchain_version_request: Option<PythonVersionRequest>,
+) -> Result<PathBuf, Error> {
     let app_dir = get_app_dir();
     let venv_dir = app_dir.join("self");
     let pip_tools_dir = app_dir.join("pip-tools");
@@ -94,12 +103,22 @@ pub fn ensure_self_venv(output: CommandOutput) -> Result<PathBuf, Error> {
         echo!("Bootstrapping rye internals");
     }
 
-    let version = ensure_self_toolchain(output).with_context(|| {
-        format!(
-            "failed to fetch internal cpython toolchain {}",
-            SELF_PYTHON_TARGET_VERSION
-        )
-    })?;
+    let version = match toolchain_version_request {
+        Some(ref version_request) => ensure_specific_self_toolchain(output, version_request)
+            .with_context(|| {
+                format!(
+                    "failed to provision internal cpython toolchain {}",
+                    version_request
+                )
+            })?,
+        None => ensure_latest_self_toolchain(output).with_context(|| {
+            format!(
+                "failed to fetch internal cpython toolchain {}",
+                SELF_PYTHON_TARGET_VERSION
+            )
+        })?,
+    };
+
     let py_bin = get_toolchain_python_bin(&version)?;
 
     // linux specific detection of shared libraries.
@@ -313,23 +332,59 @@ pub fn is_self_compatible_toolchain(version: &PythonVersion) -> bool {
     version.name == "cpython" && version.major == 3 && version.minor >= 9 && version.minor <= 12
 }
 
-fn ensure_self_toolchain(output: CommandOutput) -> Result<PythonVersion, Error> {
-    let possible_versions = list_known_toolchains()?
+/// Ensure that the toolchain for the self environment is available.
+fn ensure_latest_self_toolchain(output: CommandOutput) -> Result<PythonVersion, Error> {
+    if let Some(version) = list_known_toolchains()?
         .into_iter()
         .map(|x| x.0)
         .filter(is_self_compatible_toolchain)
-        .collect::<Vec<_>>();
-
-    if let Some(version) = possible_versions.into_iter().max() {
-        echo!(
-            "Found a compatible python version: {}",
-            style(&version).cyan()
-        );
+        .collect::<Vec<_>>()
+        .into_iter()
+        .max()
+    {
+        if output != CommandOutput::Quiet {
+            echo!(
+                "Found a compatible python version: {}",
+                style(&version).cyan()
+            );
+        }
         Ok(version)
     } else {
         fetch(&SELF_PYTHON_TARGET_VERSION, output)
     }
 }
+
+/// Ensure a specific toolchain is available.
+fn ensure_specific_self_toolchain(
+    output: CommandOutput,
+    toolchain_version_request: &PythonVersionRequest,
+) -> Result<PythonVersion, Error> {
+    let toolchain_version = latest_available_python_version(toolchain_version_request)
+        .ok_or_else(|| anyhow!("requested toolchain version is not available"))?;
+    if !is_self_compatible_toolchain(&toolchain_version) {
+        bail!(
+            "the requested toolchain version ({}) is not supported for rye-internal usage",
+            toolchain_version
+        );
+    }
+    if !get_toolchain_python_bin(&toolchain_version)?.is_file() {
+        if output != CommandOutput::Quiet {
+            echo!(
+                "Fetching requested internal toolchain '{}'",
+                toolchain_version
+            );
+        }
+        fetch(&toolchain_version.into(), output)
+    } else {
+        if output != CommandOutput::Quiet {
+            echo!(
+                "Found a compatible python version: {}",
+                style(&toolchain_version).cyan()
+            );
+        }
+        Ok(toolchain_version)
+    }
+}
 /// Fetches a version if missing.
 pub fn fetch(
     version: &PythonVersionRequest,

rye/src/cli/rye.rs

@@ -15,8 +15,8 @@ use self_replace::self_delete_outside_path;
 use tempfile::tempdir;
 
 use crate::bootstrap::{
-    download_url, download_url_ignore_404, ensure_self_venv, is_self_compatible_toolchain,
-    update_core_shims, SELF_PYTHON_TARGET_VERSION,
+    download_url, download_url_ignore_404, ensure_self_venv_with_toolchain,
+    is_self_compatible_toolchain, update_core_shims, SELF_PYTHON_TARGET_VERSION,
 };
 use crate::cli::toolchain::register_toolchain;
 use crate::config::Config;
@@ -93,6 +93,9 @@ pub struct InstallCommand {
     /// Register a specific toolchain before bootstrap.
     #[arg(long)]
     toolchain: Option<PathBuf>,
+    /// Use a specific toolchain version.
+    #[arg(long)]
+    toolchain_version: Option<PythonVersionRequest>,
 }
 
 #[derive(Debug, Copy, Clone)]
@@ -259,6 +262,7 @@ fn install(args: InstallCommand) -> Result<(), Error> {
             InstallMode::Default
         },
         args.toolchain.as_deref(),
+        args.toolchain_version,
     )
 }
 
@@ -343,7 +347,11 @@ fn is_fish() -> bool {
     Shell::infer().map_or(false, |x| matches!(x, Shell::Fish))
 }
 
-fn perform_install(mode: InstallMode, toolchain_path: Option<&Path>) -> Result<(), Error> {
+fn perform_install(
+    mode: InstallMode,
+    toolchain_path: Option<&Path>,
+    toolchain_version: Option<PythonVersionRequest>,
+) -> Result<(), Error> {
     let mut config = Config::current();
     let mut registered_toolchain: Option<PythonVersionRequest> = None;
     let config_doc = Arc::make_mut(&mut config).doc_mut();
@@ -353,6 +361,26 @@ fn perform_install(mode: InstallMode, toolchain_path: Option<&Path>) -> Result<(
     let target = shims.join("rye").with_extension(EXE_EXTENSION);
     let mut prompt_for_toolchain_later = false;
 
+    // When we perform an install and a toolchain path has not been passed,
+    // we always also pick up on the RYE_TOOLCHAIN environment variable
+    // as a fallback.
+    let toolchain_path = match toolchain_path {
+        Some(path) => Some(Cow::Borrowed(path)),
+        None => env::var_os("RYE_TOOLCHAIN")
+            .map(PathBuf::from)
+            .map(Cow::Owned),
+    };
+
+    // Also pick up the target version from the RYE_TOOLCHAIN_VERSION
+    // environment variable.
+    let toolchain_version_request = match toolchain_version {
+        Some(version) => Some(version),
+        None => match env::var("RYE_TOOLCHAIN_VERSION") {
+            Ok(val) => Some(val.parse()?),
+            Err(_) => None,
+        },
+    };
+
     echo!("{}", style("Welcome to Rye!").bold());
 
     if matches!(mode, InstallMode::AutoInstall) {
@@ -378,6 +406,18 @@ fn perform_install(mode: InstallMode, toolchain_path: Option<&Path>) -> Result<(
     echo!("{}", style("Details:").bold());
     echo!("  Rye Version: {}", style(env!("CARGO_PKG_VERSION")).cyan());
     echo!("  Platform: {} ({})", style(OS).cyan(), style(ARCH).cyan());
+    if let Some(ref toolchain_path) = toolchain_path {
+        echo!(
+            "  Internal Toolchain Path: {}",
+            style(toolchain_path.display()).cyan()
+        );
+    }
+    if let Some(ref toolchain_version_request) = toolchain_version_request {
+        echo!(
+            "  Internal Toolchain Version: {}",
+            style(toolchain_version_request).cyan()
+        );
+    }
 
     if cfg!(windows) && !symlinks_supported() {
         echo!();
@@ -423,7 +463,12 @@ fn perform_install(mode: InstallMode, toolchain_path: Option<&Path>) -> Result<(
         // can fill in the default.
         if !matches!(mode, InstallMode::NoPrompts) {
             if toolchain_path.is_none() {
-                prompt_for_default_toolchain(SELF_PYTHON_TARGET_VERSION, config_doc)?;
+                prompt_for_default_toolchain(
+                    toolchain_version_request
+                        .clone()
+                        .unwrap_or(SELF_PYTHON_TARGET_VERSION),
+                    config_doc,
+                )?;
             } else {
                 prompt_for_toolchain_later = true;
             }
@@ -457,7 +502,7 @@ fn perform_install(mode: InstallMode, toolchain_path: Option<&Path>) -> Result<(
             "Registering toolchain at {}",
             style(toolchain_path.display()).cyan()
         );
-        let version = register_toolchain(toolchain_path, None, |ver| {
+        let version = register_toolchain(&toolchain_path, None, |ver| {
             if ver.name != "cpython" {
                 bail!("Only cpython toolchains are allowed, got '{}'", ver.name);
             } else if !is_self_compatible_toolchain(ver) {
@@ -473,7 +518,8 @@ fn perform_install(mode: InstallMode, toolchain_path: Option<&Path>) -> Result<(
     }
 
     // Ensure internals next
-    let self_path = ensure_self_venv(CommandOutput::Normal)?;
+    let self_path =
+        ensure_self_venv_with_toolchain(CommandOutput::Normal, toolchain_version_request)?;
     echo!(
         "Updated self-python installation at {}",
         style(self_path.display()).cyan()
@@ -576,10 +622,6 @@ pub fn auto_self_install() -> Result<bool, Error> {
         return Ok(false);
     }
 
-    // auto install reads RYE_TOOLCHAIN to pre-register a
-    // regular toolchain.
-    let toolchain_path = env::var_os("RYE_TOOLCHAIN");
-
     let app_dir = get_app_dir();
     let rye_exe = app_dir
         .join("shims")
@@ -597,10 +639,7 @@ pub fn auto_self_install() -> Result<bool, Error> {
             crate::request_continue_prompt();
         }
 
-        perform_install(
-            InstallMode::AutoInstall,
-            toolchain_path.as_ref().map(Path::new),
-        )?;
+        perform_install(InstallMode::AutoInstall, None, None)?;
         Ok(true)
     }
 }

rye/src/cli/shim.rs

@@ -13,7 +13,7 @@ use crate::config::Config;
 use crate::consts::VENV_BIN;
 use crate::platform::{get_python_version_request_from_pyenv_pin, get_toolchain_python_bin};
 use crate::pyproject::{latest_available_python_version, PyProject};
-use crate::sources::{PythonVersion, PythonVersionRequest};
+use crate::sources::PythonVersionRequest;
 use crate::sync::{sync, SyncOptions};
 use crate::tui::redirect_to_stderr;
 use crate::utils::{exec_spawn, get_venv_python_bin, CommandOutput};
@@ -253,11 +253,8 @@ fn get_shim_target(
             return find_shadowed_target(target, args);
         };
 
-        let py_ver = match PythonVersion::try_from(version_request.clone()) {
-            Ok(py_ver) => py_ver,
-            Err(_) => latest_available_python_version(&version_request)
-                .ok_or_else(|| anyhow!("Unable to determine target Python version"))?,
-        };
+        let py_ver = latest_available_python_version(&version_request)
+            .ok_or_else(|| anyhow!("Unable to determine target Python version"))?;
         let py = get_toolchain_python_bin(&py_ver)?;
         if !py.is_file() {
             let hint = if implicit_request {