From a13421ee177af84794aa58c25ac0f12d26284c14 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 2 Jul 2026 10:11:57 -0500 Subject: [PATCH] Reduce site-packages scan allocation overhead --- crates/uv-installer/src/site_packages.rs | 86 +++++++++++++++++------- 1 file changed, 61 insertions(+), 25 deletions(-) diff --git a/crates/uv-installer/src/site_packages.rs b/crates/uv-installer/src/site_packages.rs index ccca781db4b6a..93a2c4d217c61 100644 --- a/crates/uv-installer/src/site_packages.rs +++ b/crates/uv-installer/src/site_packages.rs @@ -1,5 +1,4 @@ use std::borrow::Cow; -use std::collections::BTreeSet; use std::iter::Flatten; use std::path::PathBuf; @@ -59,30 +58,12 @@ impl SitePackages { for site_packages in interpreter.site_packages() { // Read the site-packages directory. let site_packages = match fs::read_dir(site_packages.as_ref()) { - Ok(read_dir) => { - // Collect sorted directory paths; `read_dir` is not stable across platforms - let dist_likes: BTreeSet<_> = read_dir - .filter_map(|read_dir| match read_dir { - Ok(entry) => match entry.file_type() { - Ok(file_type) => (file_type.is_dir() - || entry - .path() - .extension() - .is_some_and(|ext| ext == "egg-link" || ext == "egg-info")) - .then_some(Ok(entry.path())), - Err(err) => Some(Err(err)), - }, - Err(err) => Some(Err(err)), - }) - .collect::>() - .with_context(|| { - format!( - "Failed to read site-packages directory contents: {}", - site_packages.user_display() - ) - })?; - dist_likes - } + Ok(read_dir) => sorted_dist_like_paths(read_dir).with_context(|| { + format!( + "Failed to read site-packages directory contents: {}", + site_packages.user_display() + ) + })?, Err(err) if err.kind() == std::io::ErrorKind::NotFound => { return Ok(Self { interpreter: interpreter.clone(), @@ -624,6 +605,25 @@ impl IntoIterator for SitePackages { } } +fn sorted_dist_like_paths(read_dir: fs::ReadDir) -> Result, std::io::Error> { + let mut paths = read_dir + .filter_map(|read_dir| match read_dir { + Ok(entry) => match entry.file_type() { + Ok(file_type) => (file_type.is_dir() + || entry + .path() + .extension() + .is_some_and(|ext| ext == "egg-link" || ext == "egg-info")) + .then_some(Ok(entry.path())), + Err(err) => Some(Err(err)), + }, + Err(err) => Some(Err(err)), + }) + .collect::, std::io::Error>>()?; + paths.sort_unstable(); + Ok(paths) +} + #[derive(Debug)] pub enum SitePackagesDiagnostic { MetadataUnavailable { @@ -746,3 +746,39 @@ impl InstalledPackagesProvider for SitePackages { self.get_packages(name) } } + +#[cfg(test)] +mod tests { + use anyhow::Result; + + use super::sorted_dist_like_paths; + + #[test] + fn sorted_dist_like_paths_filters_and_sorts() -> Result<()> { + let site_packages = tempfile::tempdir()?; + fs_err::create_dir(site_packages.path().join("z_package-1.0.0.dist-info"))?; + fs_err::create_dir(site_packages.path().join("a_package"))?; + fs_err::write(site_packages.path().join("editable.egg-link"), "")?; + fs_err::write(site_packages.path().join("module.py"), "")?; + fs_err::write(site_packages.path().join("metadata.egg-info"), "")?; + + let paths = sorted_dist_like_paths(fs_err::read_dir(site_packages.path())?)?; + let names = paths + .iter() + .filter_map(|path| path.file_name()) + .map(|name| name.to_string_lossy().into_owned()) + .collect::>(); + + assert_eq!( + names, + vec![ + "a_package".to_string(), + "editable.egg-link".to_string(), + "metadata.egg-info".to_string(), + "z_package-1.0.0.dist-info".to_string(), + ] + ); + + Ok(()) + } +}