Fix mastodon-py dist-info handling (#336)

mastodon-py 1.5.1 uses a dot in its dist-info dir name, which we previously didn't handle, causing home-assistant to fail. The new implementation is based on https://github.com/pypa/packaging/blob/2f83540272e79e3fe1f5d42abae8df0c14ddf4c2/src/packaging/utils.py#L146-L172. Part of #199 ``` unzip -l Mastodon.py-1.5.1-py2.py3-none-any.whl Archive: Mastodon.py-1.5.1-py2.py3-none-any.whl Length Date Time Name --------- ---------- ----- ---- 153929 2020-02-29 17:39 mastodon/Mastodon.py 1029 2019-10-11 19:15 mastodon/__init__.py 7357 2019-10-11 20:24 mastodon/streaming.py 10 2020-03-14 18:14 Mastodon.py-1.5.1.dist-info/DESCRIPTION.rst 1398 2020-03-14 18:14 Mastodon.py-1.5.1.dist-info/metadata.json 9 2020-03-14 18:14 Mastodon.py-1.5.1.dist-info/top_level.txt 110 2020-03-14 18:14 Mastodon.py-1.5.1.dist-info/WHEEL 1543 2020-03-14 18:14 Mastodon.py-1.5.1.dist-info/METADATA 753 2020-03-14 18:14 Mastodon.py-1.5.1.dist-info/RECORD --------- ------- 166138 9 files ```
astral-sh · Nov 7, 2023 · fbe28d3 · fbe28d3
1 parent aac8ae9
commit fbe28d3
Show file tree

Hide file tree

Showing 9 changed files with 71 additions and 72 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/install-wheel-rs/Cargo.toml b/crates/install-wheel-rs/Cargo.toml
@@ -18,7 +18,9 @@ name = "install_wheel_rs"
 
 [dependencies]
 distribution-filename = { path = "../distribution-filename" }
+pep440_rs = { path = "../pep440-rs" }
 platform-host = { path = "../platform-host" }
+puffin-normalize = { path = "../puffin-normalize" }
 pypi-types = { path = "../pypi-types" }
 
 clap = { workspace = true, optional = true, features = ["derive", "env"] }

diff --git a/crates/install-wheel-rs/src/lib.rs b/crates/install-wheel-rs/src/lib.rs
@@ -1,20 +1,23 @@
 //! Takes a wheel and installs it into a venv..
 
 use std::io;
+use std::str::FromStr;
 
 use distribution_filename::WheelFilename;
 use platform_info::PlatformInfoError;
 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;
@@ -75,28 +78,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: <https://github.com/pypa/packaging/blob/2f83540272e79e3fe1f5d42abae8df0c14ddf4c2/src/packaging/utils.py#L146-L172>
+pub fn find_dist_info<'a, T: Copy>(
     filename: &WheelFilename,
     files: impl Iterator<Item = (T, &'a str)>,
 ) -> 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());
         }
@@ -106,11 +110,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::<Vec<_>>()
                     .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
@@ -25,7 +25,7 @@ use pypi_types::DirectUrl;
 use crate::install_location::{InstallLocation, LockedDir};
 use crate::record::RecordEntry;
 use crate::script::Script;
-use crate::Error;
+use crate::{find_dist_info, Error};
 
 /// `#!/usr/bin/env python`
 pub const SHEBANG_PYTHON: &str = "#!/usr/bin/env python";
@@ -930,7 +930,10 @@ 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 = find_dist_info(filename, archive.file_names().map(|name| (name, name)))
+        .map_err(Error::InvalidWheel)?
+        .1
+        .to_string();
     let (name, _version) = read_metadata(&dist_info_prefix, &mut archive)?;
     // TODO: Check that name and version match
 
@@ -1033,45 +1036,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
-///
-/// <https://github.com/PyO3/python-pkginfo-rs>
-pub fn find_dist_info(
-    filename: &WheelFilename,
-    archive: &mut ZipArchive<impl Read + Seek + Sized>,
-) -> Result<String, Error> {
-    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)
-}
-
 /// <https://github.com/PyO3/python-pkginfo-rs>
 fn read_metadata(
     dist_info_prefix: &str,

diff --git 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 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
@@ -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 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-dispatch/src/lib.rs b/crates/puffin-dispatch/src/lib.rs
@@ -60,7 +60,7 @@ impl BuildContext for BuildDispatch {
         &self.base_python
     }
 
-    #[instrument(skip(self))]
+    #[instrument(skip(self, requirements))]
     fn resolve<'a>(
         &'a self,
         requirements: &'a [Requirement],

diff --git 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 pypi_types::Metadata21;
@@ -51,10 +52,12 @@ impl CachedWheel {
     /// Read the [`Metadata21`] from a wheel.
     pub(super) fn read_dist_info(&self) -> Result<Metadata21> {
         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 dist_info = std::io::read_to_string(
-            archive.by_name(&format!("{dist_info_prefix}.dist-info/METADATA"))?,
-        )?;
+        let filename = &self.filename;
+        let dist_info_dir = find_dist_info(filename, archive.file_names().map(|name| (name, name)))
+            .map_err(|err| format_err!("Invalid wheel {filename}: {err}"))?
+            .1;
+        let dist_info =
+            std::io::read_to_string(archive.by_name(&format!("{dist_info_dir}/METADATA"))?)?;
         Ok(Metadata21::parse(dist_info.as_bytes())?)
     }
 }
diff --git a/crates/puffin-resolver/src/distribution/wheel.rs b/crates/puffin-resolver/src/distribution/wheel.rs
@@ -1,7 +1,7 @@
 use std::path::Path;
 use std::str::FromStr;
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use fs_err::tokio as fs;
 
 use tokio_util::compat::FuturesAsyncReadCompatExt;
@@ -34,7 +34,7 @@ impl<'a> WheelFetcher<'a> {
     ) -> Result<Option<Metadata21>> {
         CachedWheel::find_in_cache(distribution, tags, self.0.join(REMOTE_WHEELS_CACHE))
             .as_ref()
-            .map(CachedWheel::read_dist_info)
+            .map(|wheel| CachedWheel::read_dist_info(wheel).context("Failed to read dist info"))
             .transpose()
     }