diff --git a/Cargo.lock b/Cargo.lock index b2eefb7123e5a..392077644bc9f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1441,9 +1441,11 @@ dependencies = [ "indoc", "mailparse", "once_cell", + "pep440_rs 0.3.12", "platform-host", "platform-info", "plist", + "puffin-normalize", "pyo3", "rayon", "reflink-copy", diff --git a/crates/install-wheel-rs/Cargo.toml b/crates/install-wheel-rs/Cargo.toml index 7b02bf5773195..2ce387d80b45a 100644 --- a/crates/install-wheel-rs/Cargo.toml +++ b/crates/install-wheel-rs/Cargo.toml @@ -17,8 +17,10 @@ license = { workspace = true } name = "install_wheel_rs" [dependencies] -platform-host = { path = "../platform-host" } distribution-filename = { path = "../distribution-filename" } +pep440_rs = { path = "../pep440-rs" } +platform-host = { path = "../platform-host" } +puffin-normalize = { path = "../puffin-normalize" } clap = { workspace = true, optional = true, features = ["derive", "env"] } configparser = { workspace = true } diff --git a/crates/install-wheel-rs/src/lib.rs b/crates/install-wheel-rs/src/lib.rs index 6615d8341fea5..52e1241b0612a 100644 --- a/crates/install-wheel-rs/src/lib.rs +++ b/crates/install-wheel-rs/src/lib.rs @@ -1,6 +1,7 @@ //! Takes a wheel and installs it into a venv.. use std::io; +use std::str::FromStr; use distribution_filename::WheelFilename; use platform_info::PlatformInfoError; @@ -8,13 +9,15 @@ use thiserror::Error; use zip::result::ZipError; pub use install_location::{normalize_name, InstallLocation, LockedDir}; +use pep440_rs::Version; use platform_host::{Arch, Os}; +use puffin_normalize::PackageName; pub use record::RecordEntry; pub use script::Script; pub use uninstall::{uninstall_wheel, Uninstall}; pub use wheel::{ - find_dist_info, get_script_launcher, install_wheel, parse_key_value_file, read_record_file, - relative_to, SHEBANG_PYTHON, + get_script_launcher, install_wheel, parse_key_value_file, read_record_file, relative_to, + SHEBANG_PYTHON, }; mod install_location; @@ -73,28 +76,29 @@ impl Error { /// The metadata name may be uppercase, while the wheel and dist info names are lowercase, or /// the metadata name and the dist info name are lowercase, while the wheel name is uppercase. -/// Either way, we just search the wheel for the name -pub fn find_dist_info_metadata<'a, T: Copy>( +/// Either way, we just search the wheel for the name. +/// +/// Reference implementation: +pub fn find_dist_info<'a, T: Copy>( filename: &WheelFilename, files: impl Iterator, ) -> Result<(T, &'a str), String> { - let dist_info_matcher = format!( - "{}-{}", - filename.distribution.as_dist_info_name(), - filename.version - ); let metadatas: Vec<_> = files .filter_map(|(payload, path)| { - let (dir, file) = path.split_once('/')?; - let dir = dir.strip_suffix(".dist-info")?; - if dir.to_lowercase() == dist_info_matcher && file == "METADATA" { - Some((payload, path)) + let (dist_info_dir, file) = path.split_once('/')?; + let dir_stem = dist_info_dir.strip_suffix(".dist-info")?; + let (name, version) = dir_stem.rsplit_once('-')?; + if PackageName::from_str(name).ok()? == filename.distribution + && Version::from_str(version).ok()? == filename.version + && file == "METADATA" + { + Some((payload, dist_info_dir)) } else { None } }) .collect(); - let (payload, path) = match metadatas[..] { + let (payload, dist_info_dir) = match metadatas[..] { [] => { return Err("no .dist-info directory".to_string()); } @@ -104,11 +108,37 @@ pub fn find_dist_info_metadata<'a, T: Copy>( "multiple .dist-info directories: {}", metadatas .into_iter() - .map(|(_, path)| path.to_string()) + .map(|(_, dist_info_dir)| dist_info_dir.to_string()) .collect::>() .join(", ") )); } }; - Ok((payload, path)) + Ok((payload, dist_info_dir)) +} + +#[cfg(test)] +mod test { + use crate::find_dist_info; + use distribution_filename::WheelFilename; + use std::str::FromStr; + + #[test] + fn test_dot_in_name() { + let files = [ + "mastodon/Mastodon.py", + "mastodon/__init__.py", + "mastodon/streaming.py", + "Mastodon.py-1.5.1.dist-info/DESCRIPTION.rst", + "Mastodon.py-1.5.1.dist-info/metadata.json", + "Mastodon.py-1.5.1.dist-info/top_level.txt", + "Mastodon.py-1.5.1.dist-info/WHEEL", + "Mastodon.py-1.5.1.dist-info/METADATA", + "Mastodon.py-1.5.1.dist-info/RECORD", + ]; + let filename = WheelFilename::from_str("Mastodon.py-1.5.1-py2.py3-none-any.whl").unwrap(); + let (_, dist_info_dir) = + find_dist_info(&filename, files.into_iter().map(|file| (file, file))).unwrap(); + assert_eq!(dist_info_dir, "Mastodon.py-1.5.1.dist-info"); + } } diff --git a/crates/install-wheel-rs/src/wheel.rs b/crates/install-wheel-rs/src/wheel.rs index 384b024abbdb5..416300e9d38bd 100644 --- a/crates/install-wheel-rs/src/wheel.rs +++ b/crates/install-wheel-rs/src/wheel.rs @@ -924,7 +924,11 @@ pub fn install_wheel( ZipArchive::new(reader).map_err(|err| Error::from_zip_error("(index)".to_string(), err))?; debug!(name = name.as_ref(), "Getting wheel metadata"); - let dist_info_prefix = find_dist_info(filename, &mut archive)?; + let dist_info_prefix = + (match crate::find_dist_info(filename, archive.file_names().map(|name| (name, name))) { + Ok((_, dist_info_dir)) => Ok(dist_info_dir.to_string()), + Err(err) => Err(Error::InvalidWheel(err)), + })?; let (name, _version) = read_metadata(&dist_info_prefix, &mut archive)?; // TODO: Check that name and version match @@ -1021,45 +1025,6 @@ pub fn install_wheel( Ok(filename.get_tag()) } -/// The metadata name may be uppercase, while the wheel and dist info names are lowercase, or -/// the metadata name and the dist info name are lowercase, while the wheel name is uppercase. -/// Either way, we just search the wheel for the name -/// -/// -pub fn find_dist_info( - filename: &WheelFilename, - archive: &mut ZipArchive, -) -> Result { - let dist_info_matcher = format!( - "{}-{}", - filename.distribution.as_dist_info_name(), - filename.version - ) - .to_lowercase(); - let dist_infos: Vec<_> = archive - .file_names() - .filter_map(|name| name.split_once('/')) - .filter_map(|(dir, file)| Some((dir.strip_suffix(".dist-info")?, file))) - .filter(|(dir, file)| dir.to_lowercase() == dist_info_matcher && *file == "METADATA") - .map(|(dir, _file)| dir) - .collect(); - let dist_info = match dist_infos.as_slice() { - [] => { - return Err(Error::InvalidWheel( - "Missing .dist-info directory".to_string(), - )); - } - [dist_info] => (*dist_info).to_string(), - _ => { - return Err(Error::InvalidWheel(format!( - "Multiple .dist-info directories: {}", - dist_infos.join(", ") - ))); - } - }; - Ok(dist_info) -} - /// fn read_metadata( dist_info_prefix: &str, diff --git a/crates/puffin-client/src/client.rs b/crates/puffin-client/src/client.rs index 1c1429b467a96..237e2d5442dd8 100644 --- a/crates/puffin-client/src/client.rs +++ b/crates/puffin-client/src/client.rs @@ -19,7 +19,7 @@ use tracing::{debug, trace}; use url::Url; use distribution_filename::WheelFilename; -use install_wheel_rs::find_dist_info_metadata; +use install_wheel_rs::find_dist_info; use puffin_normalize::PackageName; use puffin_package::pypi_types::{File, Metadata21, SimpleJson}; @@ -274,16 +274,14 @@ impl RegistryClient { .await .map_err(|err| Error::Zip(filename.clone(), err))?; - let ((metadata_idx, _metadata_entry), _path) = find_dist_info_metadata( + let (metadata_idx, _dist_info_dir) = find_dist_info( filename, reader .file() .entries() .iter() .enumerate() - .filter_map(|(idx, e)| { - Some(((idx, e), e.entry().filename().as_str().ok()?)) - }), + .filter_map(|(idx, e)| Some((idx, e.entry().filename().as_str().ok()?))), ) .map_err(|err| Error::InvalidDistInfo(filename.clone(), err))?; diff --git a/crates/puffin-client/src/remote_metadata.rs b/crates/puffin-client/src/remote_metadata.rs index 516558eb618e5..205bc169a2111 100644 --- a/crates/puffin-client/src/remote_metadata.rs +++ b/crates/puffin-client/src/remote_metadata.rs @@ -8,7 +8,7 @@ use tokio_util::compat::TokioAsyncReadCompatExt; use url::Url; use distribution_filename::WheelFilename; -use install_wheel_rs::find_dist_info_metadata; +use install_wheel_rs::find_dist_info; use puffin_cache::CanonicalUrl; use puffin_package::pypi_types::Metadata21; @@ -106,7 +106,7 @@ pub(crate) async fn wheel_metadata_from_remote_zip( .await .map_err(|err| Error::Zip(filename.clone(), err))?; - let ((metadata_idx, metadata_entry), _path) = find_dist_info_metadata( + let ((metadata_idx, metadata_entry), _path) = find_dist_info( filename, reader .file() diff --git a/crates/puffin-resolver/src/distribution/cached_wheel.rs b/crates/puffin-resolver/src/distribution/cached_wheel.rs index 7036c3dce143a..7a2140104b10e 100644 --- a/crates/puffin-resolver/src/distribution/cached_wheel.rs +++ b/crates/puffin-resolver/src/distribution/cached_wheel.rs @@ -1,10 +1,11 @@ use std::path::{Path, PathBuf}; use std::str::FromStr; -use anyhow::Result; +use anyhow::{format_err, Result}; use zip::ZipArchive; use distribution_filename::WheelFilename; +use install_wheel_rs::find_dist_info; use platform_tags::Tags; use puffin_distribution::RemoteDistributionRef; use puffin_package::pypi_types::Metadata21; @@ -51,7 +52,12 @@ impl CachedWheel { /// Read the [`Metadata21`] from a wheel. pub(super) fn read_dist_info(&self) -> Result { let mut archive = ZipArchive::new(fs_err::File::open(&self.path)?)?; - let dist_info_prefix = install_wheel_rs::find_dist_info(&self.filename, &mut archive)?; + let filename = &self.filename; + let dist_info_prefix = + find_dist_info(filename, archive.file_names().map(|name| (name, name))) + .map_err(|err| format_err!("Invalid wheel {filename}: {err}"))? + .1 + .to_string(); let dist_info = std::io::read_to_string( archive.by_name(&format!("{dist_info_prefix}.dist-info/METADATA"))?, )?;