Skip to content

Commit

Permalink
Allow multiple toolchains to be requested in uv toolchain install
Browse files Browse the repository at this point in the history
  • Loading branch information
zanieb committed Jun 17, 2024
1 parent 5248269 commit 25427db
Show file tree
Hide file tree
Showing 8 changed files with 126 additions and 71 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

19 changes: 15 additions & 4 deletions crates/uv-toolchain/src/downloads.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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("-"))
}
Expand Down Expand Up @@ -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,
Expand Down
1 change: 0 additions & 1 deletion crates/uv-toolchain/src/toolchain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand Down
1 change: 1 addition & 0 deletions crates/uv/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"] }
Expand Down
4 changes: 2 additions & 2 deletions crates/uv/src/cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
163 changes: 103 additions & 60 deletions crates/uv/src/commands/toolchain/install.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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,
Expand All @@ -27,88 +29,129 @@ 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()
)?;
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);
}
} else {
unfilled_requests.push(request);
}
}

if force {
writeln!(printer.stderr(), "Forcing reinstallation...")?;
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 {
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()
)?;
}
return Ok(ExitStatus::Success);
writeln!(printer.stderr(), "Requested toolchain already installed.")?;
}
return Ok(ExitStatus::Success);
}

// Fill platform information missing from the request
let request = PythonDownloadRequest::from_request(request)?.fill()?;
let s = if unfilled_requests.len() == 1 {
""
} else {
"s"
};
writeln!(
printer.stderr(),
"Installing {} toolchain{s}",
unfilled_requests.len()
)?;

// Find the corresponding download
let download = PythonDownload::from_request(&request)?;
let version = download.python_version();
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()
.connectivity(connectivity)
.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)
}
2 changes: 1 addition & 1 deletion crates/uv/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
6 changes: 3 additions & 3 deletions crates/uv/src/settings.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}

Expand All @@ -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 }
}
}

Expand Down

0 comments on commit 25427db

Please sign in to comment.