mirrors/uv
1
0
Fork 0
mirror of https://github.com/astral-sh/uv.git synced 2025-08-04 02:48:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Refactor distribution types to adhere to a clear hierarchy (#369)

Browse source 
## Summary

This PR refactors our `RemoteDistribution` type such that it now follows
a clear hierarchy that matches the actual variants, and encodes the
differences between source and built distributions:

```rust
pub enum Distribution {
    Built(BuiltDistribution),
    Source(SourceDistribution),
}

pub enum BuiltDistribution {
    Registry(RegistryBuiltDistribution),
    DirectUrl(DirectUrlBuiltDistribution),
}

pub enum SourceDistribution {
    Registry(RegistrySourceDistribution),
    DirectUrl(DirectUrlSourceDistribution),
    Git(GitSourceDistribution),
}

/// A built distribution (wheel) that exists in a registry, like `PyPI`.
pub struct RegistryBuiltDistribution {
    pub name: PackageName,
    pub version: Version,
    pub file: File,
}

/// A built distribution (wheel) that exists at an arbitrary URL.
pub struct DirectUrlBuiltDistribution {
    pub name: PackageName,
    pub url: Url,
}

/// A source distribution that exists in a registry, like `PyPI`.
pub struct RegistrySourceDistribution {
    pub name: PackageName,
    pub version: Version,
    pub file: File,
}

/// A source distribution that exists at an arbitrary URL.
pub struct DirectUrlSourceDistribution {
    pub name: PackageName,
    pub url: Url,
}

/// A source distribution that exists in a Git repository.
pub struct GitSourceDistribution {
    pub name: PackageName,
    pub url: Url,
}
```

Most of the PR just stems downstream from this change. There are no
behavioral changes, so I'm largely relying on lint, tests, and the
compiler for correctness.
This commit is contained in:
Charlie Marsh 2023-11-09 18:45:41 -08:00 • committed by GitHub
parent 33c0901a28
commit a148f9d0be
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
36 changed files with 1729 additions and 1098 deletions
Download patch file Download diff file Expand all files Collapse all files

28
crates/puffin-resolver/src/distribution/wheel.rs → crates/puffin-resolver/src/distribution/built_distribution.rs
View file

@ -3,14 +3,13 @@ use std::str::FromStr;
use anyhow::{Context, Result};
use fs_err::tokio as fs;
use tokio_util::compat::FuturesAsyncReadCompatExt;
use tracing::debug;
use distribution_filename::WheelFilename;
use platform_tags::Tags;
use puffin_client::RegistryClient;
use puffin_distribution::RemoteDistributionRef;
use puffin_distribution::{DirectUrlBuiltDistribution, DistributionIdentifier, RemoteDistribution};
use pypi_types::Metadata21;
use crate::distribution::cached_wheel::CachedWheel;
@ -18,10 +17,10 @@ use crate::distribution::cached_wheel::CachedWheel;
const REMOTE_WHEELS_CACHE: &str = "remote-wheels-v0";
/// Fetch a built distribution from a remote source, or from a local cache.
pub(crate) struct WheelFetcher<'a>(&'a Path);
pub(crate) struct BuiltDistributionFetcher<'a>(&'a Path);
impl<'a> WheelFetcher<'a> {
/// Initialize a [`WheelFetcher`] from a [`BuildContext`].
impl<'a> BuiltDistributionFetcher<'a> {
/// Initialize a [`BuiltDistributionFetcher`] from a [`BuildContext`].
pub(crate) fn new(cache: &'a Path) -> Self {
Self(cache)
}
@ -29,7 +28,7 @@ impl<'a> WheelFetcher<'a> {
/// Read the [`Metadata21`] from a wheel, if it exists in the cache.
pub(crate) fn find_dist_info(
&self,
distribution: &RemoteDistributionRef<'_>,
distribution: &DirectUrlBuiltDistribution,
tags: &Tags,
) -> Result<Option<Metadata21>> {
CachedWheel::find_in_cache(distribution, tags, self.0.join(REMOTE_WHEELS_CACHE))
@ -41,26 +40,27 @@ impl<'a> WheelFetcher<'a> {
/// Download a wheel, storing it in the cache.
pub(crate) async fn download_wheel(
&self,
distribution: &RemoteDistributionRef<'_>,
distribution: &DirectUrlBuiltDistribution,
client: &RegistryClient,
) -> Result<Metadata21> {
debug!("Downloading: {distribution}");
let url = distribution.url()?;
let reader = client.stream_external(&url).await?;
let mut reader = tokio::io::BufReader::new(reader.compat());
let reader = client.stream_external(&distribution.url).await?;
// Create a directory for the wheel.
let wheel_dir = self.0.join(REMOTE_WHEELS_CACHE).join(distribution.id());
let wheel_dir = self
.0
.join(REMOTE_WHEELS_CACHE)
.join(distribution.distribution_id());
fs::create_dir_all(&wheel_dir).await?;
// Download the wheel.
let wheel_filename = distribution.filename()?;
let wheel_file = wheel_dir.join(wheel_filename.as_ref());
let wheel_file = wheel_dir.join(wheel_filename);
let mut writer = tokio::fs::File::create(&wheel_file).await?;
tokio::io::copy(&mut reader, &mut writer).await?;
tokio::io::copy(&mut reader.compat(), &mut writer).await?;
// Read the metadata from the wheel.
let wheel = CachedWheel::new(wheel_file, WheelFilename::from_str(&wheel_filename)?);
let wheel = CachedWheel::new(wheel_file, WheelFilename::from_str(wheel_filename)?);
let metadata21 = wheel.read_dist_info()?;
debug!("Finished downloading: {distribution}");

8
crates/puffin-resolver/src/distribution/cached_wheel.rs
View file

@ -7,7 +7,7 @@ use zip::ZipArchive;
use distribution_filename::WheelFilename;
use install_wheel_rs::find_dist_info;
use platform_tags::Tags;
use puffin_distribution::RemoteDistributionRef;
use puffin_distribution::DistributionIdentifier;
use pypi_types::Metadata21;
/// A cached wheel built from a remote source.
@ -23,12 +23,12 @@ impl CachedWheel {
}
/// Search for a wheel matching the tags that was built from the given distribution.
pub(super) fn find_in_cache(
distribution: &RemoteDistributionRef<'_>,
pub(super) fn find_in_cache<T: DistributionIdentifier>(
distribution: &T,
tags: &Tags,
cache: impl AsRef<Path>,
) -> Option<Self> {
let wheel_dir = cache.as_ref().join(distribution.id());
let wheel_dir = cache.as_ref().join(distribution.distribution_id());
let Ok(read_dir) = fs_err::read_dir(wheel_dir) else {
return None;
};

4
crates/puffin-resolver/src/distribution/mod.rs
View file

@ -1,8 +1,8 @@
pub(crate) use built_distribution::BuiltDistributionFetcher;
pub(crate) use source_distribution::{
Reporter as SourceDistributionReporter, SourceDistributionFetcher,
};
pub(crate) use wheel::WheelFetcher;
mod built_distribution;
mod cached_wheel;
mod source_distribution;
mod wheel;

99
crates/puffin-resolver/src/distribution/source_distribution.rs
View file

@ -7,8 +7,6 @@ use std::sync::Arc;
use anyhow::{bail, Result};
use fs_err::tokio as fs;
use tempfile::tempdir_in;
use tokio_util::compat::FuturesAsyncReadCompatExt;
use tracing::debug;
use url::Url;
@ -16,9 +14,9 @@ use url::Url;
use distribution_filename::WheelFilename;
use platform_tags::Tags;
use puffin_client::RegistryClient;
use puffin_distribution::source::Source;
use puffin_distribution::RemoteDistributionRef;
use puffin_git::{Git, GitSource};
use puffin_distribution::direct_url::{DirectArchiveUrl, DirectGitUrl};
use puffin_distribution::{DistributionIdentifier, RemoteDistribution, SourceDistribution};
use puffin_git::{GitSource, GitUrl};
use puffin_traits::BuildContext;
use pypi_types::Metadata21;
@ -55,7 +53,7 @@ impl<'a, T: BuildContext> SourceDistributionFetcher<'a, T> {
/// Read the [`Metadata21`] from a built source distribution, if it exists in the cache.
pub(crate) fn find_dist_info(
&self,
distribution: &RemoteDistributionRef<'_>,
distribution: &SourceDistribution,
tags: &Tags,
) -> Result<Option<Metadata21>> {
CachedWheel::find_in_cache(
@ -71,7 +69,7 @@ impl<'a, T: BuildContext> SourceDistributionFetcher<'a, T> {
/// Download and build a source distribution, storing the built wheel in the cache.
pub(crate) async fn download_and_build_sdist(
&self,
distribution: &RemoteDistributionRef<'_>,
distribution: &SourceDistribution,
client: &RegistryClient,
) -> Result<Metadata21> {
debug!("Building: {distribution}");
@ -80,50 +78,54 @@ impl<'a, T: BuildContext> SourceDistributionFetcher<'a, T> {
bail!("Building source distributions is disabled");
}
// This could extract the subdirectory.
let source = Source::try_from(distribution)?;
let (sdist_file, subdirectory) = match source {
Source::RegistryUrl(url) => {
debug!("Fetching source distribution from registry: {url}");
let (sdist_file, subdirectory) = match distribution {
SourceDistribution::Registry(sdist) => {
debug!(
"Fetching source distribution from registry: {}",
sdist.file.url
);
let url = Url::parse(&sdist.file.url)?;
let reader = client.stream_external(&url).await?;
// Download the source distribution.
let temp_dir = tempfile::tempdir_in(self.build_context.cache())?.into_path();
let sdist_filename = sdist.filename()?;
let sdist_file = temp_dir.join(sdist_filename);
let mut writer = tokio::fs::File::create(&sdist_file).await?;
tokio::io::copy(&mut reader.compat(), &mut writer).await?;
(sdist_file, None)
}
SourceDistribution::DirectUrl(sdist) => {
debug!("Fetching source distribution from URL: {}", sdist.url);
let DirectArchiveUrl { url, subdirectory } = DirectArchiveUrl::from(&sdist.url);
let reader = client.stream_external(&url).await?;
let mut reader = tokio::io::BufReader::new(reader.compat());
// Download the source distribution.
let temp_dir = tempdir_in(self.build_context.cache())?.into_path();
let sdist_filename = distribution.filename()?;
let sdist_file = temp_dir.join(sdist_filename.as_ref());
let mut writer = tokio::fs::File::create(&sdist_file).await?;
tokio::io::copy(&mut reader, &mut writer).await?;
// Registry dependencies can't specify a subdirectory.
let subdirectory = None;
(sdist_file, subdirectory)
}
Source::RemoteUrl(url, subdirectory) => {
debug!("Fetching source distribution from URL: {url}");
let reader = client.stream_external(url).await?;
let mut reader = tokio::io::BufReader::new(reader.compat());
// Download the source distribution.
let temp_dir = tempdir_in(self.build_context.cache())?.into_path();
let sdist_filename = distribution.filename()?;
let sdist_file = temp_dir.join(sdist_filename.as_ref());
let temp_dir = tempfile::tempdir_in(self.build_context.cache())?.into_path();
let sdist_filename = sdist.filename()?;
let sdist_file = temp_dir.join(sdist_filename);
let mut writer = tokio::fs::File::create(&sdist_file).await?;
tokio::io::copy(&mut reader, &mut writer).await?;
(sdist_file, subdirectory)
}
Source::Git(git, subdirectory) => {
debug!("Fetching source distribution from Git: {git}");
SourceDistribution::Git(sdist) => {
debug!("Fetching source distribution from Git: {}", sdist.url);
let DirectGitUrl { url, subdirectory } = DirectGitUrl::try_from(&sdist.url)?;
let git_dir = self.build_context.cache().join(GIT_CACHE);
let source = if let Some(reporter) = &self.reporter {
GitSource::new(git, git_dir).with_reporter(Facade::from(reporter.clone()))
GitSource::new(url, git_dir).with_reporter(Facade::from(reporter.clone()))
} else {
GitSource::new(git, git_dir)
GitSource::new(url, git_dir)
};
let sdist_file = tokio::task::spawn_blocking(move || source.fetch())
.await??
@ -138,7 +140,7 @@ impl<'a, T: BuildContext> SourceDistributionFetcher<'a, T> {
.build_context
.cache()
.join(BUILT_WHEELS_CACHE)
.join(distribution.id());
.join(distribution.distribution_id());
fs::create_dir_all(&wheel_dir).await?;
// Build the wheel.
@ -171,17 +173,15 @@ impl<'a, T: BuildContext> SourceDistributionFetcher<'a, T> {
/// This method takes into account various normalizations that are independent from the Git
/// layer. For example: removing `#subdirectory=pkg_dir`-like fragments, and removing `git+`
/// prefix kinds.
pub(crate) async fn precise(
&self,
distribution: &RemoteDistributionRef<'_>,
) -> Result<Option<Url>> {
let source = Source::try_from(distribution)?;
let Source::Git(git, subdirectory) = source else {
pub(crate) async fn precise(&self, distribution: &SourceDistribution) -> Result<Option<Url>> {
let SourceDistribution::Git(sdist) = distribution else {
return Ok(None);
};
let DirectGitUrl { url, subdirectory } = DirectGitUrl::try_from(&sdist.url)?;
// If the commit already contains a complete SHA, short-circuit.
if git.precise().is_some() {
if url.precise().is_some() {
return Ok(None);
}
@ -189,16 +189,15 @@ impl<'a, T: BuildContext> SourceDistributionFetcher<'a, T> {
// commit, etc.).
let git_dir = self.build_context.cache().join(GIT_CACHE);
let source = if let Some(reporter) = &self.reporter {
GitSource::new(git, git_dir).with_reporter(Facade::from(reporter.clone()))
GitSource::new(url, git_dir).with_reporter(Facade::from(reporter.clone()))
} else {
GitSource::new(git, git_dir)
GitSource::new(url, git_dir)
};
let precise = tokio::task::spawn_blocking(move || source.fetch()).await??;
let git = Git::from(precise);
let url = GitUrl::from(precise);
// Re-encode as a URL.
let source = Source::Git(git, subdirectory);
Ok(Some(source.into()))
Ok(Some(DirectGitUrl { url, subdirectory }.into()))
}
}

57
crates/puffin-resolver/src/error.rs
View file

@ -6,6 +6,7 @@ use thiserror::Error;
use url::Url;
use pep508_rs::Requirement;
use puffin_distribution::{BuiltDistribution, SourceDistribution};
use puffin_normalize::PackageName;
use crate::pubgrub::{PubGrubPackage, PubGrubVersion};
@ -46,16 +47,32 @@ pub enum ResolveError {
#[error("Package `{0}` attempted to resolve via URL: {1}. URL dependencies must be expressed as direct requirements or constraints. Consider adding `{0} @ {1}` to your dependencies or constraints file.")]
DisallowedUrl(PackageName, Url),
#[error("Failed to build distribution: {filename}")]
RegistryDistribution {
#[error("Failed to fetch wheel metadata from: {filename}")]
RegistryBuiltDistribution {
filename: String,
// TODO(konstin): Gives this a proper error type
#[source]
err: anyhow::Error,
},
#[error("Failed to build distribution: {url}")]
UrlDistribution {
#[error("Failed to fetch wheel metadata from: {url}")]
UrlBuiltDistribution {
url: Url,
// TODO(konstin): Gives this a proper error type
#[source]
err: anyhow::Error,
},
#[error("Failed to build distribution: {filename}")]
RegistrySourceDistribution {
filename: String,
// TODO(konstin): Gives this a proper error type
#[source]
err: anyhow::Error,
},
#[error("Failed to build distribution from URL: {url}")]
UrlSourceDistribution {
url: Url,
// TODO(konstin): Gives this a proper error type
#[source]
@ -93,3 +110,35 @@ impl From<pubgrub::error::PubGrubError<PubGrubPackage, Range<PubGrubVersion>>> f
ResolveError::PubGrub(RichPubGrubError { source: value })
}
}
impl ResolveError {
pub fn from_source_distribution(distribution: SourceDistribution, err: anyhow::Error) -> Self {
match distribution {
SourceDistribution::Registry(sdist) => Self::RegistrySourceDistribution {
filename: sdist.file.filename.clone(),
err,
},
SourceDistribution::DirectUrl(sdist) => Self::UrlSourceDistribution {
url: sdist.url.clone(),
err,
},
SourceDistribution::Git(sdist) => Self::UrlSourceDistribution {
url: sdist.url.clone(),
err,
},
}
}
pub fn from_built_distribution(distribution: BuiltDistribution, err: anyhow::Error) -> Self {
match distribution {
BuiltDistribution::Registry(wheel) => Self::RegistryBuiltDistribution {
filename: wheel.file.filename.clone(),
err,
},
BuiltDistribution::DirectUrl(wheel) => Self::UrlBuiltDistribution {
url: wheel.url.clone(),
err,
},
}
}
}

5
crates/puffin-resolver/src/file.rs
View file

@ -1,14 +1,13 @@
use distribution_filename::{SourceDistributionFilename, WheelFilename};
use std::ops::Deref;
use pypi_types::File;
/// A distribution can either be a wheel or a source distribution.
#[derive(Debug, Clone)]
pub(crate) struct WheelFile(pub(crate) File, pub(crate) WheelFilename);
pub(crate) struct WheelFile(pub(crate) File);
#[derive(Debug, Clone)]
pub(crate) struct SdistFile(pub(crate) File, pub(crate) SourceDistributionFilename);
pub(crate) struct SdistFile(pub(crate) File);
#[derive(Debug, Clone)]
pub(crate) enum DistributionFile {

19
crates/puffin-resolver/src/finder.rs
View file

@ -13,7 +13,7 @@ use distribution_filename::{SourceDistributionFilename, WheelFilename};
use pep508_rs::{Requirement, VersionOrUrl};
use platform_tags::Tags;
use puffin_client::RegistryClient;
use puffin_distribution::RemoteDistribution;
use puffin_distribution::Distribution;
use puffin_normalize::PackageName;
use pypi_types::{File, SimpleJson};
@ -66,7 +66,7 @@ impl<'a> DistributionFinder<'a> {
.ready_chunks(32);
// Resolve the requirements.
let mut resolution: FxHashMap<PackageName, RemoteDistribution> =
let mut resolution: FxHashMap<PackageName, Distribution> =
FxHashMap::with_capacity_and_hasher(requirements.len(), BuildHasherDefault::default());
// Push all the requirements into the package sink.
@ -77,7 +77,7 @@ impl<'a> DistributionFinder<'a> {
}
Some(VersionOrUrl::Url(url)) => {
let package_name = requirement.name.clone();
let package = RemoteDistribution::from_url(package_name.clone(), url.clone());
let package = Distribution::from_url(package_name.clone(), url.clone());
resolution.insert(package_name, package);
}
}
@ -126,7 +126,7 @@ impl<'a> DistributionFinder<'a> {
}
/// select a version that satisfies the requirement, preferring wheels to source distributions.
fn select(&self, requirement: &Requirement, files: Vec<File>) -> Option<RemoteDistribution> {
fn select(&self, requirement: &Requirement, files: Vec<File>) -> Option<Distribution> {
let mut fallback = None;
for file in files.into_iter().rev() {
if let Ok(wheel) = WheelFilename::from_str(file.filename.as_str()) {
@ -134,7 +134,7 @@ impl<'a> DistributionFinder<'a> {
continue;
}
if requirement.is_satisfied_by(&wheel.version) {
return Some(RemoteDistribution::from_registry(
return Some(Distribution::from_registry(
wheel.distribution,
wheel.version,
file,
@ -144,11 +144,7 @@ impl<'a> DistributionFinder<'a> {
SourceDistributionFilename::parse(file.filename.as_str(), &requirement.name)
{
if requirement.is_satisfied_by(&sdist.version) {
fallback = Some(RemoteDistribution::from_registry(
sdist.name,
sdist.version,
file,
));
fallback = Some(Distribution::from_registry(sdist.name, sdist.version, file));
}
}
}
@ -157,7 +153,6 @@ impl<'a> DistributionFinder<'a> {
}
#[derive(Debug)]
#[allow(clippy::large_enum_variant)]
enum Request {
/// A request to fetch the metadata for a package.
Package(Requirement),
@ -171,7 +166,7 @@ enum Response {
pub trait Reporter: Send + Sync {
/// Callback to invoke when a package is resolved to a specific distribution.
fn on_progress(&self, wheel: &RemoteDistribution);
fn on_progress(&self, wheel: &Distribution);
/// Callback to invoke when the resolution is complete.
fn on_complete(&self);

1
crates/puffin-resolver/src/lib.rs
View file

@ -12,6 +12,7 @@ mod distribution;
mod error;
mod file;
mod finder;
mod locks;
mod manifest;
mod prerelease_mode;
mod pubgrub;

23
crates/puffin-resolver/src/locks.rs Normal file
View file

@ -0,0 +1,23 @@
use std::sync::Arc;
use fxhash::FxHashMap;
use tokio::sync::Mutex;
use puffin_distribution::DistributionIdentifier;
/// A set of locks used to prevent concurrent access to the same resource.
#[derive(Debug, Default)]
pub(crate) struct Locks(Mutex<FxHashMap<String, Arc<Mutex<()>>>>);
impl Locks {
/// Acquire a lock on the given resource.
pub(crate) async fn acquire(
&self,
distribution: &impl DistributionIdentifier,
) -> Arc<Mutex<()>> {
let mut map = self.0.lock().await;
map.entry(distribution.resource_id())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
}
}

51
crates/puffin-resolver/src/resolution.rs
View file

@ -11,7 +11,7 @@ use waitmap::WaitMap;
use pep440_rs::{Version, VersionSpecifier, VersionSpecifiers};
use pep508_rs::{Requirement, VersionOrUrl};
use puffin_distribution::RemoteDistribution;
use puffin_distribution::{BaseDistribution, BuiltDistribution, Distribution, SourceDistribution};
use puffin_normalize::PackageName;
use pypi_types::File;
@ -19,21 +19,21 @@ use crate::pubgrub::{PubGrubPackage, PubGrubPriority, PubGrubVersion};
/// A set of packages pinned at specific versions.
#[derive(Debug, Default)]
pub struct Resolution(FxHashMap<PackageName, RemoteDistribution>);
pub struct Resolution(FxHashMap<PackageName, Distribution>);
impl Resolution {
/// Create a new resolution from the given pinned packages.
pub(crate) fn new(packages: FxHashMap<PackageName, RemoteDistribution>) -> Self {
pub(crate) fn new(packages: FxHashMap<PackageName, Distribution>) -> Self {
Self(packages)
}
/// Return the distribution for the given package name, if it exists.
pub fn get(&self, package_name: &PackageName) -> Option<&RemoteDistribution> {
pub fn get(&self, package_name: &PackageName) -> Option<&Distribution> {
self.0.get(package_name)
}
/// Iterate over the [`RemoteDistribution`] entities in this resolution.
pub fn into_distributions(self) -> impl Iterator<Item = RemoteDistribution> {
/// Iterate over the [`Distribution`] entities in this resolution.
pub fn into_distributions(self) -> impl Iterator<Item = Distribution> {
self.0.into_values()
}
@ -51,7 +51,7 @@ impl Resolution {
/// A complete resolution graph in which every node represents a pinned package and every edge
/// represents a dependency between two pinned packages.
#[derive(Debug)]
pub struct Graph(petgraph::graph::Graph<RemoteDistribution, (), petgraph::Directed>);
pub struct Graph(petgraph::graph::Graph<Distribution, (), petgraph::Directed>);
impl Graph {
/// Create a new graph from the resolved `PubGrub` state.
@ -78,7 +78,7 @@ impl Graph {
.unwrap()
.clone();
let pinned_package =
RemoteDistribution::from_registry(package_name.clone(), version, file);
Distribution::from_registry(package_name.clone(), version, file);
let index = graph.add_node(pinned_package);
inverse.insert(package_name, index);
@ -87,8 +87,7 @@ impl Graph {
let url = redirects
.get(url)
.map_or_else(|| url.clone(), |url| url.value().clone());
let pinned_package =
RemoteDistribution::from_url(package_name.clone(), url.clone());
let pinned_package = Distribution::from_url(package_name.clone(), url);
let index = graph.add_node(pinned_package);
inverse.insert(package_name, index);
@ -144,18 +143,38 @@ impl Graph {
self.0
.node_indices()
.map(|node| match &self.0[node] {
RemoteDistribution::Registry(name, version, _file) => Requirement {
name: name.clone(),
Distribution::Built(BuiltDistribution::Registry(wheel)) => Requirement {
name: wheel.name.clone(),
extras: None,
version_or_url: Some(VersionOrUrl::VersionSpecifier(VersionSpecifiers::from(
VersionSpecifier::equals_version(version.clone()),
VersionSpecifier::equals_version(wheel.version.clone()),
))),
marker: None,
},
RemoteDistribution::Url(name, url) => Requirement {
name: name.clone(),
Distribution::Built(BuiltDistribution::DirectUrl(wheel)) => Requirement {
name: wheel.name.clone(),
extras: None,
version_or_url: Some(VersionOrUrl::Url(url.clone())),
version_or_url: Some(VersionOrUrl::Url(wheel.url.clone())),
marker: None,
},
Distribution::Source(SourceDistribution::Registry(sdist)) => Requirement {
name: sdist.name.clone(),
extras: None,
version_or_url: Some(VersionOrUrl::VersionSpecifier(VersionSpecifiers::from(
VersionSpecifier::equals_version(sdist.version.clone()),
))),
marker: None,
},
Distribution::Source(SourceDistribution::DirectUrl(sdist)) => Requirement {
name: sdist.name.clone(),
extras: None,
version_or_url: Some(VersionOrUrl::Url(sdist.url.clone())),
marker: None,
},
Distribution::Source(SourceDistribution::Git(sdist)) => Requirement {
name: sdist.name.clone(),
extras: None,
version_or_url: Some(VersionOrUrl::Url(sdist.url.clone())),
marker: None,
},
})

434
crates/puffin-resolver/src/resolver.rs
View file

@ -13,7 +13,6 @@ use pubgrub::range::Range;
use pubgrub::solver::{Incompatibility, State};
use pubgrub::type_aliases::DependencyConstraints;
use tokio::select;
use tokio::sync::Mutex;
use tracing::{debug, error, trace};
use url::Url;
use waitmap::WaitMap;
@ -21,17 +20,23 @@ use waitmap::WaitMap;
use distribution_filename::{SourceDistributionFilename, WheelFilename};
use pep508_rs::{MarkerEnvironment, Requirement};
use platform_tags::Tags;
use puffin_cache::{CanonicalUrl, RepositoryUrl};
use puffin_cache::CanonicalUrl;
use puffin_client::RegistryClient;
use puffin_distribution::{RemoteDistributionRef, VersionOrUrl};
use puffin_distribution::{
BaseDistribution, BuiltDistribution, DirectUrlSourceDistribution, Distribution,
DistributionIdentifier, GitSourceDistribution, SourceDistribution, VersionOrUrl,
};
use puffin_normalize::{ExtraName, PackageName};
use puffin_traits::BuildContext;
use pypi_types::{File, Metadata21, SimpleJson};
use crate::candidate_selector::CandidateSelector;
use crate::distribution::{SourceDistributionFetcher, SourceDistributionReporter, WheelFetcher};
use crate::distribution::{
BuiltDistributionFetcher, SourceDistributionFetcher, SourceDistributionReporter,
};
use crate::error::ResolveError;
use crate::file::{DistributionFile, SdistFile, WheelFile};
use crate::locks::Locks;
use crate::manifest::Manifest;
use crate::pubgrub::{
PubGrubDependencies, PubGrubPackage, PubGrubPriorities, PubGrubVersion, MIN_VERSION,
@ -285,18 +290,11 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
}
}
PubGrubPackage::Package(package_name, _extra, Some(url)) => {
// Emit a request to fetch the metadata for this package.
// Emit a request to fetch the metadata for this distribution.
if in_flight.insert_url(url) {
priorities.add(package_name.clone());
if WheelFilename::try_from(url).is_ok() {
// Kick off a request to download the wheel.
request_sink
.unbounded_send(Request::WheelUrl(package_name.clone(), url.clone()))?;
} else {
// Otherwise, assume this is a source distribution.
request_sink
.unbounded_send(Request::SdistUrl(package_name.clone(), url.clone()))?;
}
let distribution = Distribution::from_url(package_name.clone(), url.clone());
request_sink.unbounded_send(Request::Distribution(distribution))?;
}
}
}
@ -332,24 +330,13 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
};
// Emit a request to fetch the metadata for this version.
match candidate.file {
DistributionFile::Wheel(file) => {
if in_flight.insert_file(&file) {
request_sink.unbounded_send(Request::Wheel(
candidate.package_name.clone(),
file.clone(),
))?;
}
}
DistributionFile::Sdist(file) => {
if in_flight.insert_file(&file) {
request_sink.unbounded_send(Request::Sdist(
candidate.package_name.clone(),
candidate.version.clone().into(),
file.clone(),
))?;
}
}
if in_flight.insert_file(&candidate.file) {
let distribution = Distribution::from_registry(
candidate.package_name.clone(),
candidate.version.clone().into(),
candidate.file.clone().into(),
);
request_sink.unbounded_send(Request::Distribution(distribution))?;
}
}
Ok(())
@ -389,7 +376,12 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
}
} else {
// Otherwise, assume this is a source distribution.
let entry = self.index.versions.wait(url.as_str()).await.unwrap();
let entry = self
.index
.distributions
.wait(&url.distribution_id())
.await
.unwrap();
let metadata = entry.value();
let version = PubGrubVersion::from(metadata.version.clone());
if range.contains(&version) {
@ -430,24 +422,13 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
);
// Emit a request to fetch the metadata for this version.
match candidate.file {
DistributionFile::Wheel(file) => {
if in_flight.insert_file(&file) {
request_sink.unbounded_send(Request::Wheel(
candidate.package_name.clone(),
file.clone(),
))?;
}
}
DistributionFile::Sdist(file) => {
if in_flight.insert_file(&file) {
request_sink.unbounded_send(Request::Sdist(
candidate.package_name.clone(),
candidate.version.clone().into(),
file.clone(),
))?;
}
}
if in_flight.insert_file(&candidate.file) {
let distribution = Distribution::from_registry(
candidate.package_name.clone(),
candidate.version.clone().into(),
candidate.file.clone().into(),
);
request_sink.unbounded_send(Request::Distribution(distribution))?;
}
let version = candidate.version.clone();
@ -490,11 +471,20 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
PubGrubPackage::Package(package_name, extra, url) => {
// Wait for the metadata to be available.
let entry = match url {
Some(url) => self.index.versions.wait(url.as_str()).await.unwrap(),
Some(url) => self
.index
.distributions
.wait(&url.distribution_id())
.await
.unwrap(),
None => {
let versions = pins.get(package_name).unwrap();
let file = versions.get(version.into()).unwrap();
self.index.versions.wait(&file.hashes.sha256).await.unwrap()
self.index
.distributions
.wait(&file.distribution_id())
.await
.unwrap()
}
};
let metadata = entry.value();
@ -555,15 +545,11 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
std::collections::btree_map::Entry::Occupied(mut entry) => {
if matches!(entry.get(), DistributionFile::Sdist(_)) {
// Wheels get precedence over source distributions.
entry.insert(DistributionFile::from(WheelFile(
file, filename,
)));
entry.insert(DistributionFile::from(WheelFile(file)));
}
}
std::collections::btree_map::Entry::Vacant(entry) => {
entry.insert(DistributionFile::from(WheelFile(
file, filename,
)));
entry.insert(DistributionFile::from(WheelFile(file)));
}
}
}
@ -574,7 +560,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
if let std::collections::btree_map::Entry::Vacant(entry) =
version_map.entry(version)
{
entry.insert(DistributionFile::from(SdistFile(file, filename)));
entry.insert(DistributionFile::from(SdistFile(file)));
}
}
}
@ -583,30 +569,27 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
.packages
.insert(package_name.clone(), version_map);
}
Response::Wheel(file, metadata) => {
trace!("Received wheel metadata for: {}", file.filename);
Response::Distribution(Distribution::Built(distribution), metadata, ..) => {
trace!("Received built distribution metadata for: {distribution}");
self.index
.versions
.insert(file.hashes.sha256.clone(), metadata);
.distributions
.insert(distribution.distribution_id(), metadata);
}
Response::Sdist(file, metadata) => {
trace!("Received sdist metadata for: {}", file.filename);
Response::Distribution(Distribution::Source(distribution), metadata, precise) => {
trace!("Received source distribution metadata for: {distribution}");
self.index
.versions
.insert(file.hashes.sha256.clone(), metadata);
}
Response::WheelUrl(url, precise, metadata) => {
trace!("Received remote wheel metadata for: {url}");
self.index.versions.insert(url.to_string(), metadata);
.distributions
.insert(distribution.distribution_id(), metadata);
if let Some(precise) = precise {
self.index.redirects.insert(url, precise);
}
}
Response::SdistUrl(url, precise, metadata) => {
trace!("Received remote source distribution metadata for: {url}");
self.index.versions.insert(url.to_string(), metadata);
if let Some(precise) = precise {
self.index.redirects.insert(url, precise);
match distribution {
SourceDistribution::DirectUrl(sdist) => {
self.index.redirects.insert(sdist.url.clone(), precise);
}
SourceDistribution::Git(sdist) => {
self.index.redirects.insert(sdist.url.clone(), precise);
}
SourceDistribution::Registry(_) => {}
}
}
}
}
@ -625,63 +608,67 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
.map_err(ResolveError::Client)
.await
}
// Fetch wheel metadata from the registry.
Request::Wheel(package_name, file) => {
let metadata = self
.client
.wheel_metadata(file.0.clone(), file.1.clone())
.map_err(ResolveError::Client)
.await?;
if metadata.name != package_name {
// Fetch wheel metadata.
Request::Distribution(Distribution::Built(distribution)) => {
let metadata =
match &distribution {
BuiltDistribution::Registry(wheel) => {
self.client
.wheel_metadata(wheel.file.clone())
.map_err(ResolveError::Client)
.await?
}
BuiltDistribution::DirectUrl(wheel) => {
let fetcher = BuiltDistributionFetcher::new(self.build_context.cache());
match fetcher.find_dist_info(wheel, self.tags) {
Ok(Some(metadata)) => {
debug!("Found wheel metadata in cache: {wheel}");
metadata
}
Ok(None) => {
debug!("Downloading wheel: {wheel}");
fetcher.download_wheel(wheel, self.client).await.map_err(
|err| {
ResolveError::from_built_distribution(
distribution.clone(),
err,
)
},
)?
}
Err(err) => {
error!("Failed to read wheel from cache: {err}");
fetcher.download_wheel(wheel, self.client).await.map_err(
|err| {
ResolveError::from_built_distribution(
distribution.clone(),
err,
)
},
)?
}
}
}
};
if metadata.name != *distribution.name() {
return Err(ResolveError::NameMismatch {
metadata: metadata.name,
given: package_name,
given: distribution.name().clone(),
});
}
Ok(Response::Wheel(file, metadata))
Ok(Response::Distribution(
Distribution::Built(distribution),
metadata,
None,
))
}
// Build a source distribution from the registry, returning its metadata.
Request::Sdist(package_name, version, file) => {
let builder = SourceDistributionFetcher::new(self.build_context);
let distribution =
RemoteDistributionRef::from_registry(&package_name, &version, &file);
let metadata = match builder.find_dist_info(&distribution, self.tags) {
Ok(Some(metadata)) => metadata,
Ok(None) => builder
.download_and_build_sdist(&distribution, self.client)
.await
.map_err(|err| ResolveError::RegistryDistribution {
filename: file.filename.clone(),
err,
})?,
Err(err) => {
error!(
"Failed to read source distribution {distribution} from cache: {err}",
);
builder
.download_and_build_sdist(&distribution, self.client)
.await
.map_err(|err| ResolveError::RegistryDistribution {
filename: file.filename.clone(),
err,
})?
}
};
if metadata.name != package_name {
return Err(ResolveError::NameMismatch {
metadata: metadata.name,
given: package_name,
});
}
Ok(Response::Sdist(file, metadata))
}
// Build a source distribution from a remote URL, returning its metadata.
Request::SdistUrl(package_name, url) => {
let lock = self.locks.acquire(&url).await;
// Fetch source distribution metadata.
Request::Distribution(Distribution::Source(sdist)) => {
let lock = self.locks.acquire(&sdist).await;
let _guard = lock.lock().await;
let fetcher = if let Some(reporter) = &self.reporter {
@ -693,109 +680,77 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
};
let precise = fetcher
.precise(&RemoteDistributionRef::from_url(&package_name, &url))
.precise(&sdist)
.await
.map_err(|err| ResolveError::UrlDistribution {
url: url.clone(),
err,
})?;
let distribution = RemoteDistributionRef::from_url(
&package_name,
precise.as_ref().unwrap_or(&url),
);
.map_err(|err| ResolveError::from_source_distribution(sdist.clone(), err))?;
let task = self
.reporter
.as_ref()
.map(|reporter| reporter.on_build_start(&distribution));
.map(|reporter| reporter.on_build_start(&sdist));
let metadata = match fetcher.find_dist_info(&distribution, self.tags) {
Ok(Some(metadata)) => {
debug!("Found source distribution metadata in cache: {url}");
metadata
}
Ok(None) => {
debug!("Downloading source distribution from: {url}");
fetcher
.download_and_build_sdist(&distribution, self.client)
.await
.map_err(|err| ResolveError::UrlDistribution {
url: url.clone(),
err,
})?
}
Err(err) => {
error!(
"Failed to read source distribution {distribution} from cache: {err}",
);
fetcher
.download_and_build_sdist(&distribution, self.client)
.await
.map_err(|err| ResolveError::UrlDistribution {
url: url.clone(),
err,
})?
let metadata = {
// Insert the `precise`, if it exists.
let sdist = match sdist.clone() {
SourceDistribution::DirectUrl(sdist) => {
SourceDistribution::DirectUrl(DirectUrlSourceDistribution {
url: precise.clone().unwrap_or_else(|| sdist.url.clone()),
..sdist
})
}
SourceDistribution::Git(sdist) => {
SourceDistribution::Git(GitSourceDistribution {
url: precise.clone().unwrap_or_else(|| sdist.url.clone()),
..sdist
})
}
sdist @ SourceDistribution::Registry(_) => sdist,
};
match fetcher.find_dist_info(&sdist, self.tags) {
Ok(Some(metadata)) => {
debug!("Found source distribution metadata in cache: {sdist}");
metadata
}
Ok(None) => {
debug!("Downloading source distribution: {sdist}");
fetcher
.download_and_build_sdist(&sdist, self.client)
.await
.map_err(|err| {
ResolveError::from_source_distribution(sdist.clone(), err)
})?
}
Err(err) => {
error!("Failed to read source distribution from cache: {err}",);
fetcher
.download_and_build_sdist(&sdist, self.client)
.await
.map_err(|err| {
ResolveError::from_source_distribution(sdist.clone(), err)
})?
}
}
};
if metadata.name != package_name {
if metadata.name != *sdist.name() {
return Err(ResolveError::NameMismatch {
metadata: metadata.name,
given: package_name,
given: sdist.name().clone(),
});
}
if let Some(task) = task {
if let Some(reporter) = self.reporter.as_ref() {
reporter.on_build_complete(&distribution, task);
reporter.on_build_complete(&sdist, task);
}
}
Ok(Response::SdistUrl(url, precise, metadata))
}
// Fetch wheel metadata from a remote URL.
Request::WheelUrl(package_name, url) => {
let lock = self.locks.acquire(&url).await;
let _guard = lock.lock().await;
let fetcher = WheelFetcher::new(self.build_context.cache());
let distribution = RemoteDistributionRef::from_url(&package_name, &url);
let metadata = match fetcher.find_dist_info(&distribution, self.tags) {
Ok(Some(metadata)) => {
debug!("Found wheel metadata in cache: {url}");
metadata
}
Ok(None) => {
debug!("Downloading wheel from: {url}");
fetcher
.download_wheel(&distribution, self.client)
.await
.map_err(|err| ResolveError::UrlDistribution {
url: url.clone(),
err,
})?
}
Err(err) => {
error!("Failed to read wheel {distribution} from cache: {err}",);
fetcher
.download_wheel(&distribution, self.client)
.await
.map_err(|err| ResolveError::UrlDistribution {
url: url.clone(),
err,
})?
}
};
if metadata.name != package_name {
return Err(ResolveError::NameMismatch {
metadata: metadata.name,
given: package_name,
});
}
Ok(Response::WheelUrl(url, None, metadata))
Ok(Response::Distribution(
Distribution::Source(sdist),
metadata,
precise,
))
}
}
}
@ -835,10 +790,10 @@ pub trait Reporter: Send + Sync {
fn on_complete(&self);
/// Callback to invoke when a source distribution build is kicked off.
fn on_build_start(&self, distribution: &RemoteDistributionRef<'_>) -> usize;
fn on_build_start(&self, distribution: &SourceDistribution) -> usize;
/// Callback to invoke when a source distribution build is complete.
fn on_build_complete(&self, distribution: &RemoteDistributionRef<'_>, id: usize);
fn on_build_complete(&self, distribution: &SourceDistribution, id: usize);
/// Callback to invoke when a repository checkout begins.
fn on_checkout_start(&self, url: &Url, rev: &str) -> usize;
@ -864,31 +819,21 @@ impl SourceDistributionReporter for Facade {
/// Fetch the metadata for an item
#[derive(Debug)]
#[allow(clippy::large_enum_variant)]
enum Request {
/// A request to fetch the metadata for a package.
Package(PackageName),
/// A request to fetch wheel metadata from a registry.
Wheel(PackageName, WheelFile),
/// A request to fetch source distribution metadata from a registry.
Sdist(PackageName, pep440_rs::Version, SdistFile),
/// A request to fetch wheel metadata from a remote URL.
WheelUrl(PackageName, Url),
/// A request to fetch source distribution metadata from a remote URL.
SdistUrl(PackageName, Url),
/// A request to fetch the metadata for a built or source distribution.
Distribution(Distribution),
}
#[derive(Debug)]
#[allow(clippy::large_enum_variant)]
enum Response {
/// The returned metadata for a package hosted on a registry.
Package(PackageName, SimpleJson),
/// The returned metadata for a wheel hosted on a registry.
Wheel(WheelFile, Metadata21),
/// The returned metadata for a source distribution hosted on a registry.
Sdist(SdistFile, Metadata21),
/// The returned metadata for a wheel hosted on a remote URL.
WheelUrl(Url, Option<Url>, Metadata21),
/// The returned metadata for a source distribution hosted on a remote URL.
SdistUrl(Url, Option<Url>, Metadata21),
/// The returned metadata for a distribution.
Distribution(Distribution, Metadata21, Option<Url>),
}
pub(crate) type VersionMap = BTreeMap<PubGrubVersion, DistributionFile>;
@ -910,8 +855,11 @@ impl InFlight {
self.packages.insert(package_name.clone())
}
fn insert_file(&mut self, file: &File) -> bool {
self.files.insert(file.hashes.sha256.clone())
fn insert_file(&mut self, file: &DistributionFile) -> bool {
match file {
DistributionFile::Wheel(file) => self.files.insert(file.hashes.sha256.clone()),
DistributionFile::Sdist(file) => self.files.insert(file.hashes.sha256.clone()),
}
}
fn insert_url(&mut self, url: &Url) -> bool {
@ -919,27 +867,13 @@ impl InFlight {
}
}
/// A set of locks used to prevent concurrent access to the same resource.
#[derive(Debug, Default)]
struct Locks(Mutex<FxHashMap<String, Arc<Mutex<()>>>>);
impl Locks {
/// Acquire a lock on the given resource.
async fn acquire(&self, url: &Url) -> Arc<Mutex<()>> {
let mut map = self.0.lock().await;
map.entry(puffin_cache::digest(&RepositoryUrl::new(url)))
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
}
}
/// In-memory index of package metadata.
struct Index {
/// A map from package name to the metadata for that package.
packages: WaitMap<PackageName, VersionMap>,
/// A map from wheel SHA or URL to the metadata for that wheel.
versions: WaitMap<String, Metadata21>,
/// A map from distribution SHA to metadata for that distribution.
distributions: WaitMap<String, Metadata21>,
/// A map from source URL to precise URL.
redirects: WaitMap<Url, Url>,
@ -949,7 +883,7 @@ impl Default for Index {
fn default() -> Self {
Self {
packages: WaitMap::new(),
versions: WaitMap::new(),
distributions: WaitMap::new(),
redirects: WaitMap::new(),
}
}