diff --git a/crates/distribution-types/src/cached.rs b/crates/distribution-types/src/cached.rs index 456bfa77d..b3a9e891f 100644 --- a/crates/distribution-types/src/cached.rs +++ b/crates/distribution-types/src/cached.rs @@ -1,15 +1,15 @@ use std::path::{Path, PathBuf}; use std::str::FromStr; -use anyhow::{anyhow, Result}; +use anyhow::Result; use url::Url; -use crate::traits::Metadata; -use crate::{BuiltDist, Dist, SourceDist, VersionOrUrl}; -use pep440_rs::Version; +use distribution_filename::WheelFilename; use puffin_normalize::PackageName; use crate::direct_url::DirectUrl; +use crate::traits::Metadata; +use crate::{BuiltDist, Dist, SourceDist, VersionOrUrl}; /// A built distribution (wheel) that exists in the local cache. #[derive(Debug, Clone)] @@ -22,31 +22,30 @@ pub enum CachedDist { #[derive(Debug, Clone)] pub struct CachedRegistryDist { - pub name: PackageName, - pub version: Version, + pub filename: WheelFilename, pub path: PathBuf, } #[derive(Debug, Clone)] pub struct CachedDirectUrlDist { - pub name: PackageName, + pub filename: WheelFilename, pub url: Url, pub path: PathBuf, } impl Metadata for CachedRegistryDist { fn name(&self) -> &PackageName { - &self.name + &self.filename.name } fn version_or_url(&self) -> VersionOrUrl { - VersionOrUrl::Version(&self.version) + VersionOrUrl::Version(&self.filename.version) } } impl Metadata for CachedDirectUrlDist { fn name(&self) -> &PackageName { - &self.name + &self.filename.name } fn version_or_url(&self) -> VersionOrUrl { @@ -72,40 +71,36 @@ impl Metadata for CachedDist { impl CachedDist { /// Initialize a [`CachedDist`] from a [`Dist`]. - pub fn from_remote(remote: Dist, path: PathBuf) -> Self { + pub fn from_remote(remote: Dist, filename: WheelFilename, path: PathBuf) -> Self { match remote { - Dist::Built(BuiltDist::Registry(dist)) => Self::Registry(CachedRegistryDist { - name: dist.name, - version: dist.version, - path, - }), + Dist::Built(BuiltDist::Registry(_dist)) => { + Self::Registry(CachedRegistryDist { filename, path }) + } Dist::Built(BuiltDist::DirectUrl(dist)) => Self::Url(CachedDirectUrlDist { - name: dist.filename.name, + filename, url: dist.url, path, }), Dist::Built(BuiltDist::Path(dist)) => Self::Url(CachedDirectUrlDist { - name: dist.filename.name, + filename, url: dist.url, path, }), - Dist::Source(SourceDist::Registry(dist)) => Self::Registry(CachedRegistryDist { - name: dist.name, - version: dist.version, - path, - }), + Dist::Source(SourceDist::Registry(_dist)) => { + Self::Registry(CachedRegistryDist { filename, path }) + } Dist::Source(SourceDist::DirectUrl(dist)) => Self::Url(CachedDirectUrlDist { - name: dist.name, + filename, url: dist.url, path, }), Dist::Source(SourceDist::Git(dist)) => Self::Url(CachedDirectUrlDist { - name: dist.name, + filename, url: dist.url, path, }), Dist::Source(SourceDist::Path(dist)) => Self::Url(CachedDirectUrlDist { - name: dist.name, + filename, url: dist.url, path, }), @@ -130,8 +125,12 @@ impl CachedDist { } impl CachedDirectUrlDist { - pub fn from_url(name: PackageName, url: Url, path: PathBuf) -> Self { - Self { name, url, path } + pub fn from_url(filename: WheelFilename, url: Url, path: PathBuf) -> Self { + Self { + filename, + url, + path, + } } } @@ -144,18 +143,12 @@ impl CachedRegistryDist { let Some(file_name) = file_name.to_str() else { return Ok(None); }; - let Some((name, version)) = file_name.rsplit_once('-') else { + let Ok(filename) = WheelFilename::from_str(file_name) else { return Ok(None); }; - let name = PackageName::from_str(name)?; - let version = Version::from_str(version).map_err(|err| anyhow!(err))?; let path = path.to_path_buf(); - Ok(Some(Self { - name, - version, - path, - })) + Ok(Some(Self { filename, path })) } } diff --git a/crates/puffin-cli/src/commands/pip_sync.rs b/crates/puffin-cli/src/commands/pip_sync.rs index 330c792ce..d423e46bd 100644 --- a/crates/puffin-cli/src/commands/pip_sync.rs +++ b/crates/puffin-cli/src/commands/pip_sync.rs @@ -76,6 +76,8 @@ pub(crate) async fn sync_requirements( "Using Python interpreter: {}", venv.python_executable().display() ); + // Determine the current environment markers. + let tags = Tags::from_interpreter(venv.interpreter())?; // Partition into those that should be linked from the cache (`local`), those that need to be // downloaded (`remote`), and those that should be removed (`extraneous`). @@ -83,7 +85,8 @@ pub(crate) async fn sync_requirements( local, remote, extraneous, - } = InstallPlan::try_from_requirements(requirements, cache, &venv)?; + } = InstallPlan::try_from_requirements(requirements, cache, &venv, &tags) + .context("Failed to determine installation plan")?; // Nothing to do. if remote.is_empty() && local.is_empty() && extraneous.is_empty() { @@ -102,9 +105,6 @@ pub(crate) async fn sync_requirements( return Ok(ExitStatus::Success); } - // Determine the current environment markers. - let tags = Tags::from_interpreter(venv.interpreter())?; - // Instantiate a client. let client = { let mut builder = RegistryClientBuilder::new(cache); diff --git a/crates/puffin-dispatch/src/lib.rs b/crates/puffin-dispatch/src/lib.rs index 4395ce0ee..a61adcb75 100644 --- a/crates/puffin-dispatch/src/lib.rs +++ b/crates/puffin-dispatch/src/lib.rs @@ -115,13 +115,13 @@ impl BuildContext for BuildDispatch { venv.root().display(), ); + let tags = Tags::from_interpreter(&self.interpreter)?; + let InstallPlan { local, remote, extraneous, - } = InstallPlan::try_from_requirements(requirements, &self.cache, venv)?; - - let tags = Tags::from_interpreter(&self.interpreter)?; + } = InstallPlan::try_from_requirements(requirements, &self.cache, venv, &tags)?; // Resolve the dependencies. let remote = if remote.is_empty() { diff --git a/crates/puffin-installer/src/cache.rs b/crates/puffin-installer/src/cache.rs index 20819d2f3..27c2bcbfd 100644 --- a/crates/puffin-installer/src/cache.rs +++ b/crates/puffin-installer/src/cache.rs @@ -2,7 +2,9 @@ use std::path::{Path, PathBuf}; use fs_err as fs; -use distribution_types::{BuiltDist, Dist, Metadata, SourceDist}; +use distribution_filename::WheelFilename; +use distribution_types::{BuiltDist, Dist, Metadata, SourceDist, VersionOrUrl}; +use puffin_cache::{digest, CanonicalUrl}; static WHEEL_CACHE: &str = "wheels-v0"; @@ -25,10 +27,13 @@ impl WheelCache { } /// Return the path at which a given [`Dist`] would be stored. - pub(crate) fn entry(&self, dist: &Dist) -> PathBuf { - self.root - .join(CacheShard::from(dist).segment()) - .join(dist.package_id()) + pub(crate) fn entry(&self, dist: &Dist, filename: &WheelFilename) -> PathBuf { + let mut path = self.root.join(CacheShard::from(dist).segment()); + // TODO(konstin): Use `WheelMetadataCache` instead + if let VersionOrUrl::Url(url) = dist.version_or_url() { + path.push(digest(&CanonicalUrl::new(url))); + } + path.join(filename.to_string()) } /// Returns a handle to the wheel cache directory. diff --git a/crates/puffin-installer/src/plan.rs b/crates/puffin-installer/src/plan.rs index c7a488cfe..475e29977 100644 --- a/crates/puffin-installer/src/plan.rs +++ b/crates/puffin-installer/src/plan.rs @@ -1,11 +1,14 @@ use std::path::Path; +use std::str::FromStr; -use anyhow::{Context, Result}; +use anyhow::{bail, Context, Result}; +use distribution_filename::WheelFilename; use tracing::debug; use distribution_types::direct_url::DirectUrl; -use distribution_types::{CachedDist, InstalledDist}; +use distribution_types::{CachedDist, InstalledDist, RemoteSource}; use pep508_rs::{Requirement, VersionOrUrl}; +use platform_tags::Tags; use puffin_interpreter::Virtualenv; use crate::url_index::UrlIndex; @@ -33,13 +36,14 @@ impl InstallPlan { requirements: &[Requirement], cache: &Path, venv: &Virtualenv, + tags: &Tags, ) -> Result { // Index all the already-installed packages in site-packages. let mut site_packages = SitePackages::try_from_executable(venv).context("Failed to list installed packages")?; // Index all the already-downloaded wheels in the cache. - let registry_index = RegistryIndex::try_from_directory(cache); + let registry_index = RegistryIndex::try_from_directory(cache, tags); let url_index = UrlIndex::try_from_directory(cache); let mut local = vec![]; @@ -86,7 +90,7 @@ impl InstallPlan { None | Some(VersionOrUrl::VersionSpecifier(_)) => { if let Some(distribution) = registry_index .get(&requirement.name) - .filter(|dist| requirement.is_satisfied_by(&dist.version)) + .filter(|dist| requirement.is_satisfied_by(&dist.filename.version)) { debug!("Requirement already cached: {distribution}"); local.push(CachedDist::Registry(distribution.clone())); @@ -94,10 +98,24 @@ impl InstallPlan { } } Some(VersionOrUrl::Url(url)) => { - if let Some(distribution) = url_index.get(&requirement.name, url) { - debug!("Requirement already cached: {distribution}"); - local.push(CachedDist::Url(distribution.clone())); - continue; + // Only consider wheel urls + if let Some(filename) = url + .filename() + .ok() + .and_then(|filename| WheelFilename::from_str(filename).ok()) + { + if requirement.name != filename.name { + bail!( + "Given name `{}` does not match url name `{}`", + requirement.name, + url + ); + } + if let Some(distribution) = url_index.get(filename, url) { + debug!("Requirement already cached: {distribution}"); + local.push(CachedDist::Url(distribution.clone())); + continue; + } } } } diff --git a/crates/puffin-installer/src/registry_index.rs b/crates/puffin-installer/src/registry_index.rs index 5c0160e9b..4d20fb65c 100644 --- a/crates/puffin-installer/src/registry_index.rs +++ b/crates/puffin-installer/src/registry_index.rs @@ -5,6 +5,7 @@ use fs_err as fs; use tracing::warn; use distribution_types::{CachedRegistryDist, Metadata}; +use platform_tags::Tags; use puffin_normalize::PackageName; use crate::cache::{CacheShard, WheelCache}; @@ -15,7 +16,7 @@ pub struct RegistryIndex(HashMap); impl RegistryIndex { /// Build an index of cached distributions from a directory. - pub fn try_from_directory(path: &Path) -> Self { + pub fn try_from_directory(path: &Path, tags: &Tags) -> Self { let mut index = HashMap::new(); let cache = WheelCache::new(path); @@ -36,21 +37,32 @@ impl RegistryIndex { continue; } }; - if file_type.is_dir() { - match CachedRegistryDist::try_from_path(&path) { - Ok(None) => {} - Ok(Some(dist_info)) => { + if !file_type.is_dir() { + continue; + } + + match CachedRegistryDist::try_from_path(&path) { + Ok(None) => {} + Ok(Some(dist_info)) => { + // Pick the wheel with the highest priority + let compatibility = dist_info.filename.compatibility(tags); + if let Some(existing) = index.get_mut(dist_info.name()) { + // Override if we have better compatibility + if compatibility > existing.filename.compatibility(tags) { + *existing = dist_info; + } + } else if compatibility.is_some() { index.insert(dist_info.name().clone(), dist_info); } - Err(err) => { - warn!("Invalid cache entry at {}, removing. {err}", path.display()); - let result = fs::remove_dir_all(&path); - if let Err(err) = result { - warn!( - "Failed to remove invalid cache entry at {}: {err}", - path.display() - ); - } + } + Err(err) => { + warn!("Invalid cache entry at {}, removing. {err}", path.display()); + let result = fs::remove_dir_all(&path); + if let Err(err) = result { + warn!( + "Failed to remove invalid cache entry at {}: {err}", + path.display() + ); } } } diff --git a/crates/puffin-installer/src/unzipper.rs b/crates/puffin-installer/src/unzipper.rs index 3a358ad4b..40fa80576 100644 --- a/crates/puffin-installer/src/unzipper.rs +++ b/crates/puffin-installer/src/unzipper.rs @@ -44,6 +44,7 @@ impl Unzipper { let mut wheels = Vec::with_capacity(downloads.len()); for download in downloads { let remote = download.remote().clone(); + let filename = download.filename().clone(); debug!("Unpacking wheel: {remote}"); @@ -55,7 +56,7 @@ impl Unzipper { .await??; // Write the unzipped wheel to the target directory. - let target = wheel_cache.entry(&remote); + let target = wheel_cache.entry(&remote, &filename); if let Some(parent) = target.parent() { fs_err::create_dir_all(parent)?; } @@ -65,7 +66,7 @@ impl Unzipper { if let Err(err) = result { // If the renaming failed because another instance was faster, that's fine // (`DirectoryNotEmpty` is not stable so we can't match on it) - if !wheel_cache.entry(&remote).is_dir() { + if !wheel_cache.entry(&remote, &filename).is_dir() { return Err(err.into()); } } @@ -74,8 +75,8 @@ impl Unzipper { reporter.on_unzip_progress(&remote); } - let path = wheel_cache.entry(&remote); - wheels.push(CachedDist::from_remote(remote, path)); + let path = wheel_cache.entry(&remote, &filename); + wheels.push(CachedDist::from_remote(remote, filename, path)); } if let Some(reporter) = self.reporter.as_ref() { diff --git a/crates/puffin-installer/src/url_index.rs b/crates/puffin-installer/src/url_index.rs index e5b6206c3..baa4969df 100644 --- a/crates/puffin-installer/src/url_index.rs +++ b/crates/puffin-installer/src/url_index.rs @@ -4,8 +4,8 @@ use fxhash::FxHashMap; use tracing::warn; use url::Url; +use distribution_filename::WheelFilename; use distribution_types::{CachedDirectUrlDist, Identifier}; -use puffin_normalize::PackageName; use crate::cache::{CacheShard, WheelCache}; @@ -49,13 +49,13 @@ impl UrlIndex { } /// Returns a distribution from the index, if it exists. - pub(crate) fn get(&self, name: &PackageName, url: &Url) -> Option { + pub(crate) fn get(&self, filename: WheelFilename, url: &Url) -> Option { // TODO(charlie): This takes advantage of the fact that for URL dependencies, the package ID // and distribution ID are identical. We should either change the cache layout to use // distribution IDs, or implement package ID for URL. let path = self.0.get(&url.distribution_id())?; Some(CachedDirectUrlDist::from_url( - name.clone(), + filename, url.clone(), path.clone(), ))