Allow multiple toolchains to be requested in uv toolchain install (#4334)

Allows installation of multiple toolchains in a single invocation
because I don't want to be limited to one! Most of the implementation
for concurrent downloads ported from `cargo dev fetch-python`.

Cargo.lock

@@ -4434,6 +4434,7 @@ dependencies = [
  "filetime",
  "flate2",
  "fs-err",
+ "futures",
  "ignore",
  "indicatif",
  "indoc",

crates/uv-toolchain/src/downloads.rs

@@ -15,7 +15,7 @@ use uv_client::BetterReqwestError;
 use futures::TryStreamExt;
 
 use tokio_util::compat::FuturesAsyncReadCompatExt;
-use tracing::debug;
+use tracing::{debug, instrument};
 use url::Url;
 use uv_fs::Simplified;
 
@@ -265,20 +265,30 @@ impl PythonDownloadRequest {
 impl Display for PythonDownloadRequest {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let mut parts = Vec::new();
-        if let Some(version) = &self.version {
-            parts.push(version.to_string());
-        }
         if let Some(implementation) = self.implementation {
             parts.push(implementation.to_string());
+        } else {
+            parts.push("any".to_string());
+        }
+        if let Some(version) = &self.version {
+            parts.push(version.to_string());
+        } else {
+            parts.push("any".to_string());
         }
         if let Some(os) = &self.os {
             parts.push(os.to_string());
+        } else {
+            parts.push("any".to_string());
         }
         if let Some(arch) = self.arch {
             parts.push(arch.to_string());
+        } else {
+            parts.push("any".to_string());
         }
         if let Some(libc) = self.libc {
             parts.push(libc.to_string());
+        } else {
+            parts.push("any".to_string());
         }
         write!(f, "{}", parts.join("-"))
     }
@@ -371,6 +381,7 @@ impl PythonDownload {
     }
 
     /// Download and extract
+    #[instrument(skip(client, parent_path), fields(download = %self.key()))]
     pub async fn fetch(
         &self,
         client: &uv_client::BaseClient,

crates/uv-toolchain/src/toolchain.rs

@@ -82,7 +82,6 @@ impl Toolchain {
     ) -> Result<Self, Error> {
         let sources = ToolchainSources::from_settings(system, preview);
         let toolchain = find_toolchain(request, system, &sources, cache)??;
-
         Ok(toolchain)
     }
 

crates/uv/Cargo.toml

@@ -46,6 +46,7 @@ clap = { workspace = true, features = ["derive", "string", "wrap_help"] }
 clap_complete_command = { workspace = true }
 flate2 = { workspace = true, default-features = false }
 fs-err = { workspace = true, features = ["tokio"] }
+futures = { workspace = true }
 indicatif = { workspace = true }
 itertools = { workspace = true }
 miette = { workspace = true, features = ["fancy"] }

crates/uv/src/cli.rs

@@ -1754,10 +1754,10 @@ pub(crate) struct ToolchainListArgs {
 #[derive(Args)]
 #[allow(clippy::struct_excessive_bools)]
 pub(crate) struct ToolchainInstallArgs {
-    /// The toolchain to install.
+    /// The toolchains to install.
     ///
     /// If not provided, the latest available version will be installed unless a toolchain was previously installed.
-    pub(crate) target: Option<String>,
+    pub(crate) targets: Vec<String>,
 
     /// Force the installation of the toolchain, even if it is already installed.
     #[arg(long, short)]

crates/uv/src/commands/toolchain/install.rs

@@ -1,4 +1,6 @@
-use anyhow::Result;
+use anyhow::{Error, Result};
+use futures::StreamExt;
+use itertools::Itertools;
 use std::fmt::Write;
 use uv_cache::Cache;
 use uv_client::Connectivity;
@@ -15,7 +17,7 @@ use crate::printer::Printer;
 /// Download and install a Python toolchain.
 #[allow(clippy::too_many_arguments)]
 pub(crate) async fn install(
-    target: Option<String>,
+    targets: Vec<String>,
     force: bool,
     native_tls: bool,
     connectivity: Connectivity,
@@ -27,63 +29,75 @@ pub(crate) async fn install(
         warn_user!("`uv toolchain install` is experimental and may change without warning.");
     }
 
+    let start = std::time::Instant::now();
+
     let toolchains = InstalledToolchains::from_settings()?.init()?;
     let toolchain_dir = toolchains.root();
 
-    let request = if let Some(target) = target {
-        let request = ToolchainRequest::parse(&target);
-        match request {
-            ToolchainRequest::Any => (),
-            ToolchainRequest::Directory(_)
-            | ToolchainRequest::ExecutableName(_)
-            | ToolchainRequest::File(_) => {
-                writeln!(printer.stderr(), "Invalid toolchain request '{target}'")?;
-                return Ok(ExitStatus::Failure);
-            }
-            _ => {
-                writeln!(printer.stderr(), "Looking for {request}")?;
-            }
-        }
-        request
+    let requests = if targets.is_empty() {
+        vec![PythonDownloadRequest::default()]
     } else {
-        ToolchainRequest::default()
+        targets
+            .iter()
+            .map(|target| parse_target(target, printer))
+            .collect::<Result<_>>()?
     };
 
-    if let Some(toolchain) = toolchains
-        .find_all()?
-        .find(|toolchain| toolchain.satisfies(&request))
-    {
-        writeln!(
-            printer.stderr(),
-            "Found installed toolchain '{}'",
-            toolchain.key()
-        )?;
-
-        if force {
-            writeln!(printer.stderr(), "Forcing reinstallation...")?;
-        } else {
-            if matches!(request, ToolchainRequest::Any) {
-                writeln!(
-                    printer.stderr(),
-                    "A toolchain is already installed. Use `uv toolchain install <request>` to install a specific toolchain.",
-                )?;
-            } else {
-                writeln!(
-                    printer.stderr(),
-                    "Already installed at {}",
-                    toolchain.path().user_display()
-                )?;
+    let installed_toolchains: Vec<_> = toolchains.find_all()?.collect();
+    let mut unfilled_requests = Vec::new();
+    for request in requests {
+        if let Some(toolchain) = installed_toolchains
+            .iter()
+            .find(|toolchain| request.satisfied_by_key(toolchain.key()))
+        {
+            writeln!(
+                printer.stderr(),
+                "Found installed toolchain '{}' that satisfies {request}",
+                toolchain.key()
+            )?;
+            if force {
+                unfilled_requests.push(request);
             }
-            return Ok(ExitStatus::Success);
+        } else {
+            unfilled_requests.push(request);
         }
     }
 
-    // Fill platform information missing from the request
-    let request = PythonDownloadRequest::from_request(request)?.fill()?;
+    if unfilled_requests.is_empty() {
+        if targets.is_empty() {
+            writeln!(
+                printer.stderr(),
+                "A toolchain is already installed. Use `uv toolchain install <request>` to install a specific toolchain.",
+            )?;
+        } else if targets.len() > 1 {
+            writeln!(
+                printer.stderr(),
+                "All requested toolchains already installed."
+            )?;
+        } else {
+            writeln!(printer.stderr(), "Requested toolchain already installed.")?;
+        }
+        return Ok(ExitStatus::Success);
+    }
 
-    // Find the corresponding download
-    let download = PythonDownload::from_request(&request)?;
-    let version = download.python_version();
+    let s = if unfilled_requests.len() == 1 {
+        ""
+    } else {
+        "s"
+    };
+    writeln!(
+        printer.stderr(),
+        "Installing {} toolchain{s}",
+        unfilled_requests.len()
+    )?;
+
+    let downloads = unfilled_requests
+        .into_iter()
+        // Populate the download requests with defaults
+        .map(PythonDownloadRequest::fill)
+        .map(|request| request.map(|inner| PythonDownload::from_request(&inner)))
+        .flatten_ok()
+        .collect::<Result<Vec<_>, uv_toolchain::downloads::Error>>()?;
 
     // Construct a client
     let client = uv_client::BaseClientBuilder::new()
@@ -91,24 +105,53 @@ pub(crate) async fn install(
         .native_tls(native_tls)
         .build();
 
-    writeln!(printer.stderr(), "Downloading {}", download.key())?;
-    let result = download.fetch(&client, toolchain_dir).await?;
+    let mut tasks = futures::stream::iter(downloads.iter())
+        .map(|download| async {
+            let _ = writeln!(printer.stderr(), "Downloading {}", download.key());
+            let result = download.fetch(&client, toolchain_dir).await;
+            (download.python_version(), result)
+        })
+        .buffered(4);
 
-    let path = match result {
-        // Note we should only encounter `AlreadyAvailable` if there's a race condition
-        // TODO(zanieb): We should lock the toolchain directory on fetch
-        DownloadResult::AlreadyAvailable(path) => path,
-        DownloadResult::Fetched(path) => path,
-    };
-
-    let installed = InstalledToolchain::new(path)?;
-    installed.ensure_externally_managed()?;
+    let mut results = Vec::new();
+    while let Some(task) = tasks.next().await {
+        let (version, result) = task;
+        let path = match result? {
+            // We should only encounter already-available during concurrent installs
+            DownloadResult::AlreadyAvailable(path) => path,
+            DownloadResult::Fetched(path) => {
+                writeln!(
+                    printer.stderr(),
+                    "Installed Python {version} to {}",
+                    path.user_display()
+                )?;
+                path
+            }
+        };
+        // Ensure the installations have externally managed markers
+        let installed = InstalledToolchain::new(path.clone())?;
+        installed.ensure_externally_managed()?;
+        results.push((version, path));
+    }
 
+    let s = if downloads.len() == 1 { "" } else { "s" };
     writeln!(
         printer.stderr(),
-        "Installed Python {version} to {}",
-        installed.path().user_display()
+        "Installed {} toolchain{s} in {}s",
+        downloads.len(),
+        start.elapsed().as_secs()
     )?;
 
     Ok(ExitStatus::Success)
 }
+
+fn parse_target(target: &str, printer: Printer) -> Result<PythonDownloadRequest, Error> {
+    let request = ToolchainRequest::parse(target);
+    let download_request = PythonDownloadRequest::from_request(request.clone())?;
+    // TODO(zanieb): Can we improve the `PythonDownloadRequest` display?
+    writeln!(
+        printer.stderr(),
+        "Looking for toolchain {request} ({download_request})"
+    )?;
+    Ok(download_request)
+}

crates/uv/src/main.rs

@@ -796,7 +796,7 @@ async fn run() -> Result<ExitStatus> {
             let cache = cache.init()?;
 
             commands::toolchain_install(
-                args.target,
+                args.targets,
                 args.force,
                 globals.native_tls,
                 globals.connectivity,

crates/uv/src/settings.rs

@@ -257,7 +257,7 @@ impl ToolchainListSettings {
 #[allow(clippy::struct_excessive_bools)]
 #[derive(Debug, Clone)]
 pub(crate) struct ToolchainInstallSettings {
-    pub(crate) target: Option<String>,
+    pub(crate) targets: Vec<String>,
     pub(crate) force: bool,
 }
 
@@ -268,9 +268,9 @@ impl ToolchainInstallSettings {
         args: ToolchainInstallArgs,
         _filesystem: Option<FilesystemOptions>,
     ) -> Self {
-        let ToolchainInstallArgs { target, force } = args;
+        let ToolchainInstallArgs { targets, force } = args;
 
-        Self { target, force }
+        Self { targets, force }
     }
 }