diff --git a/crates/uv-python/src/downloads.rs b/crates/uv-python/src/downloads.rs index ac1bb5ed4a1d4..b278e5c31f4f8 100644 --- a/crates/uv-python/src/downloads.rs +++ b/crates/uv-python/src/downloads.rs @@ -423,7 +423,7 @@ pub enum DownloadResult { } impl ManagedPythonDownload { - /// Return the first [`PythonDownload`] matching a request, if any. + /// Return the first [`ManagedPythonDownload`] matching a request, if any. pub fn from_request( request: &PythonDownloadRequest, ) -> Result<&'static ManagedPythonDownload, Error> { @@ -433,7 +433,7 @@ impl ManagedPythonDownload { .ok_or(Error::NoDownloadFound(request.clone())) } - /// Iterate over all [`PythonDownload`]'s. + /// Iterate over all [`ManagedPythonDownload`]s. pub fn iter_all() -> impl Iterator { PYTHON_DOWNLOADS .iter() diff --git a/crates/uv/src/commands/python/install.rs b/crates/uv/src/commands/python/install.rs index dc7d2b956a9a5..30dc03397f38e 100644 --- a/crates/uv/src/commands/python/install.rs +++ b/crates/uv/src/commands/python/install.rs @@ -1,16 +1,15 @@ -use std::collections::BTreeSet; use std::fmt::Write; use std::io::ErrorKind; -use std::path::Path; +use std::path::{Path, PathBuf}; use anyhow::Result; use futures::stream::FuturesUnordered; use futures::StreamExt; -use itertools::Itertools; +use itertools::{Either, Itertools}; use owo_colors::OwoColorize; -use rustc_hash::FxHashSet; +use rustc_hash::{FxHashMap, FxHashSet}; use same_file::is_same_file; -use tracing::debug; +use tracing::{debug, trace}; use uv_client::Connectivity; use uv_configuration::TrustedHost; @@ -19,7 +18,7 @@ use uv_python::downloads::{DownloadResult, ManagedPythonDownload, PythonDownload use uv_python::managed::{ python_executable_dir, ManagedPythonInstallation, ManagedPythonInstallations, }; -use uv_python::{PythonDownloads, PythonRequest, PythonVersionFile}; +use uv_python::{PythonDownloads, PythonInstallationKey, PythonRequest, PythonVersionFile}; use uv_shell::Shell; use uv_warnings::warn_user; @@ -28,6 +27,36 @@ use crate::commands::reporters::PythonDownloadReporter; use crate::commands::{elapsed, ExitStatus}; use crate::printer::Printer; +#[derive(Debug, Clone)] +struct InstallRequest { + python: PythonRequest, + download: PythonDownloadRequest, +} + +impl InstallRequest { + fn from_python_request(request: PythonRequest) -> Result { + let download = PythonDownloadRequest::from_request(&request).ok_or_else(|| { + anyhow::anyhow!("Cannot download managed Python for request: {request}") + })?; + + Ok(Self { + python: request, + download, + }) + } + + /// Return a concrete [`ManagedPythonDownload`] from the request. + /// + /// Fills the request with platform information and finds a matching download. + fn into_download(self) -> Result<&'static ManagedPythonDownload> { + Ok(ManagedPythonDownload::from_request(&self.download.fill()?)?) + } + + fn matches_installation(&self, installation: &ManagedPythonInstallation) -> bool { + self.download.satisfied_by_key(installation.key()) + } +} + /// Download and install Python versions. #[allow(clippy::fn_params_excessive_bools)] pub(crate) async fn install( @@ -49,81 +78,77 @@ pub(crate) async fn install( let cache_dir = installations.cache(); let _lock = installations.lock().await?; - let targets = targets.into_iter().collect::>(); + let mut is_default_install = false; let requests: Vec<_> = if targets.is_empty() { PythonVersionFile::discover(project_dir, no_config, true) .await? .map(PythonVersionFile::into_versions) - .unwrap_or_else(|| vec![PythonRequest::Default]) + .unwrap_or_else(|| { + // If no version file is found and no requests were made + is_default_install = true; + vec![PythonRequest::Default] + }) + .into_iter() + .map(InstallRequest::from_python_request) + .collect::>>()? } else { targets .iter() .map(|target| PythonRequest::parse(target.as_str())) - .collect() + .map(InstallRequest::from_python_request) + .collect::>>()? }; - let download_requests = requests - .iter() - .map(|request| { - PythonDownloadRequest::from_request(request).ok_or_else(|| { - anyhow::anyhow!("Cannot download managed Python for request: {request}") - }) - }) - .collect::>>()?; - - let installed_installations: Vec<_> = installations + let existing_installations: Vec<_> = installations .find_all()? - .inspect(|installation| debug!("Found existing installation {}", installation.key())) + .inspect(|installation| trace!("Found existing installation {}", installation.key())) .collect(); - let mut unfilled_requests = Vec::new(); + + let mut existing = FxHashSet::default(); + let mut installed = FxHashSet::default(); let mut uninstalled = FxHashSet::default(); - for (request, download_request) in requests.iter().zip(download_requests) { - if matches!(requests.as_slice(), [PythonRequest::Default]) { - writeln!(printer.stderr(), "Searching for Python installations")?; - } else { - writeln!( - printer.stderr(), - "Searching for Python versions matching: {}", - request.cyan() - )?; - } - if let Some(installation) = installed_installations + + let Some(first_request) = requests.first() else { + // Nothing to do + return Ok(ExitStatus::Success); + }; + + // Find requests that are already satisfied + let (satisfied, unsatisfied): (Vec<_>, Vec<_>) = requests.iter().partition_map(|request| { + if let Some(installation) = existing_installations .iter() - .find(|installation| download_request.satisfied_by_key(installation.key())) + .find(|installation| request.matches_installation(installation)) { - if matches!(request, PythonRequest::Default) { - writeln!(printer.stderr(), "Found: {}", installation.key().green())?; + existing.insert(installation.key()); + if reinstall { + debug!( + "Ignoring match `{}` for request `{}` due to `--reinstall` flag", + installation.key().green(), + request.python.cyan() + ); + + Either::Right(request) } else { - writeln!( - printer.stderr(), - "Found existing installation for {}: {}", - request.cyan(), + debug!( + "Found `{}` for request `{}`", installation.key().green(), - )?; - } - // TODO(zanieb): Ensure executables are linked for already-installed versions - if reinstall { - uninstalled.insert(installation.key()); - unfilled_requests.push(download_request); + request.python.cyan(), + ); + + Either::Left((request, installation)) } } else { - unfilled_requests.push(download_request); - } - } + debug!( + "No installation found for request `{}`", + request.python.cyan(), + ); - if unfilled_requests.is_empty() { - if matches!(requests.as_slice(), [PythonRequest::Default]) { - writeln!( - printer.stderr(), - "Python is already available. Use `uv python install ` to install a specific version.", - )?; - } else if requests.len() > 1 { - writeln!(printer.stderr(), "All requested versions already installed")?; + Either::Right(request) } - return Ok(ExitStatus::Success); - } + }); - if matches!(python_downloads, PythonDownloads::Never) { + // Check if Python downloads are banned + if matches!(python_downloads, PythonDownloads::Never) && !unsatisfied.is_empty() { writeln!( printer.stderr(), "Python downloads are not allowed (`python-downloads = \"never\"`). Change to `python-downloads = \"manual\"` to allow explicit installs.", @@ -131,14 +156,20 @@ pub(crate) async fn install( return Ok(ExitStatus::Failure); } - let downloads = unfilled_requests - .into_iter() - // Populate the download requests with defaults - .map(|request| ManagedPythonDownload::from_request(&PythonDownloadRequest::fill(request)?)) - .collect::, uv_python::downloads::Error>>()?; - - // Ensure we only download each version once - let downloads = downloads + // Find downloads for the requests + let downloads = unsatisfied + .iter() + .map(|request| { + (*request).clone().into_download().inspect(|download| { + debug!( + "Found download `{}` for request `{}`", + download, + request.python.cyan(), + ); + }) + }) + .collect::>>()? + // Ensure we only download each version once .into_iter() .unique_by(|download| download.key()) .collect::>(); @@ -152,6 +183,7 @@ pub(crate) async fn install( let reporter = PythonDownloadReporter::new(printer, downloads.len() as u64); + // Download and unpack the Python versions concurrently let mut tasks = FuturesUnordered::new(); for download in &downloads { tasks.push(async { @@ -170,10 +202,9 @@ pub(crate) async fn install( }); } - let bin = python_executable_dir()?; - - let mut installed = FxHashSet::default(); let mut errors = vec![]; + let mut downloaded = Vec::with_capacity(unsatisfied.len()); + while let Some((key, result)) = tasks.next().await { match result { Ok(download) => { @@ -183,60 +214,113 @@ pub(crate) async fn install( DownloadResult::Fetched(path) => path, }; - installed.insert(key); + let installation = ManagedPythonInstallation::new(path)?; + installed.insert(installation.key().clone()); + if existing.contains(installation.key()) { + uninstalled.insert(installation.key().clone()); + } + downloaded.push(installation); + } + Err(err) => { + errors.push((key, anyhow::Error::new(err))); + } + } + } - // Ensure the installations have externally managed markers - let managed = ManagedPythonInstallation::new(path.clone())?; - managed.ensure_externally_managed()?; - managed.ensure_canonical_executables()?; + let bin = python_executable_dir()?; + let mut installed_executables: FxHashMap> = + FxHashMap::default(); - // TODO(zanieb): Only apply `default` for the _first_ requested version - let targets = if default { - vec![ - managed.key().executable_name_minor(), - managed.key().executable_name_major(), - managed.key().executable_name(), - ] - } else { - vec![managed.key().executable_name_minor()] - }; + // Ensure that the installations are _complete_ for both downloaded installations and existing + // installations that were requested + for installation in downloaded.iter().chain( + satisfied + .iter() + .map(|(_, installation)| installation) + .copied(), + ) { + installation.ensure_externally_managed()?; + installation.ensure_canonical_executables()?; - for target in targets { - let target = bin.join(target); - match managed.create_bin_link(&target) { - Ok(()) => { - debug!("Installed {} executable to {}", key, target.display()); + let targets = if (default || is_default_install) + && first_request.matches_installation(installation) + { + vec![ + installation.key().executable_name_minor(), + installation.key().executable_name_major(), + installation.key().executable_name(), + ] + } else { + vec![installation.key().executable_name_minor()] + }; + + for target in targets { + let target = bin.join(target); + match installation.create_bin_link(&target) { + Ok(()) => { + debug!( + "Installed executable at {} for {}", + target.user_display(), + installation.key(), + ); + installed.insert(installation.key().clone()); + match installed_executables.get_mut(installation.key()) { + Some(targets) => targets.push(target.clone()), + None => { + installed_executables + .insert(installation.key().clone(), vec![target.clone()]); } - Err(uv_python::managed::Error::LinkExecutable { from, to, err }) - if err.kind() == ErrorKind::AlreadyExists => - { - // TODO(zanieb): Add `--force` - if reinstall { - fs_err::remove_file(&to)?; - managed.create_bin_link(&target)?; - debug!("Replaced {} executable at {}", key, target.user_display()); - } else { - if !is_same_file(&to, &from).unwrap_or_default() { - errors.push(( - key, - anyhow::anyhow!( - "Executable already exists at `{}`. Use `--reinstall` to force replacement.", - to.user_display() - ), - )); - } + } + } + Err(uv_python::managed::Error::LinkExecutable { from, to, err }) + if err.kind() == ErrorKind::AlreadyExists => + { + // TODO(zanieb): Add `--force` + if reinstall { + fs_err::remove_file(&to)?; + installation.create_bin_link(&target)?; + debug!( + "Updated executable at {} to {}", + target.user_display(), + installation.key(), + ); + installed.insert(installation.key().clone()); + match installed_executables.get_mut(installation.key()) { + Some(targets) => targets.push(target.clone()), + None => { + installed_executables + .insert(installation.key().clone(), vec![target.clone()]); } } - Err(err) => return Err(err.into()), + } else { + if !is_same_file(&to, &from).unwrap_or_default() { + errors.push(( + installation.key(), + anyhow::anyhow!( + "Executable already exists at `{}`. Use `--reinstall` to force replacement.", + to.user_display() + ), + )); + } } } - } - Err(err) => { - errors.push((key, anyhow::Error::new(err))); + Err(err) => return Err(err.into()), } } } + if installed.is_empty() { + if is_default_install { + writeln!( + printer.stderr(), + "Python is already installed. Use `uv python install ` to install another version.", + )?; + } else if requests.len() > 1 { + writeln!(printer.stderr(), "All requested versions already installed")?; + } + return Ok(ExitStatus::Success); + } + if !installed.is_empty() { if installed.len() == 1 { let installed = installed.iter().next().unwrap(); @@ -267,10 +351,10 @@ pub(crate) async fn install( let reinstalled = uninstalled .intersection(&installed) - .copied() + .cloned() .collect::>(); - let uninstalled = uninstalled.difference(&reinstalled).copied(); - let installed = installed.difference(&reinstalled).copied(); + let uninstalled = uninstalled.difference(&reinstalled).cloned(); + let installed = installed.difference(&reinstalled).cloned(); for event in uninstalled .map(|key| ChangeEvent { @@ -281,7 +365,7 @@ pub(crate) async fn install( key: key.clone(), kind: ChangeEventKind::Added, })) - .chain(reinstalled.iter().map(|&key| ChangeEvent { + .chain(reinstalled.iter().map(|key| ChangeEvent { key: key.clone(), kind: ChangeEventKind::Reinstalled, })) @@ -291,28 +375,28 @@ pub(crate) async fn install( ChangeEventKind::Added => { writeln!( printer.stderr(), - " {} {} ({})", + " {} {}{}", "+".green(), event.key.bold(), - event.key.executable_name_minor() + format_installed_executables(&event.key, &installed_executables) )?; } ChangeEventKind::Removed => { writeln!( printer.stderr(), - " {} {} ({})", + " {} {}{}", "-".red(), event.key.bold(), - event.key.executable_name_minor() + format_installed_executables(&event.key, &installed_executables) )?; } ChangeEventKind::Reinstalled => { writeln!( printer.stderr(), - " {} {} ({})", + " {} {}{}", "~".yellow(), event.key.bold(), - event.key.executable_name_minor() + format_installed_executables(&event.key, &installed_executables) )?; } } @@ -344,6 +428,22 @@ pub(crate) async fn install( Ok(ExitStatus::Success) } +fn format_installed_executables( + key: &PythonInstallationKey, + installed_executables: &FxHashMap>, +) -> String { + if let Some(executables) = installed_executables.get(key) { + let executables = executables + .iter() + .filter_map(|path| path.file_name()) + .map(|name| name.to_string_lossy()) + .join(", "); + format!(" ({executables})") + } else { + String::new() + } +} + fn warn_if_not_on_path(bin: &Path) { if !Shell::contains_path(bin) { if let Some(shell) = Shell::from_env() {