From 5ee0ec447a00ff8e26158d878bb414b6c01778bf Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 1 Apr 2024 20:03:34 -0400 Subject: [PATCH] Use distribution database and index for all pre-resolution phases --- crates/distribution-types/src/buildable.rs | 10 +++ crates/uv-requirements/src/lookahead.rs | 6 +- crates/uv-requirements/src/unnamed.rs | 96 +++++++++++++--------- crates/uv/src/commands/pip_compile.rs | 13 ++- crates/uv/src/commands/pip_install.rs | 9 +- crates/uv/src/commands/pip_sync.rs | 9 +- 6 files changed, 89 insertions(+), 54 deletions(-) diff --git a/crates/distribution-types/src/buildable.rs b/crates/distribution-types/src/buildable.rs index 828c427c204e..6c1ad9a7c0cd 100644 --- a/crates/distribution-types/src/buildable.rs +++ b/crates/distribution-types/src/buildable.rs @@ -54,6 +54,16 @@ pub enum SourceUrl<'a> { Path(PathSourceUrl<'a>), } +impl<'a> SourceUrl<'a> { + pub fn url(&self) -> &Url { + match self { + Self::Direct(dist) => dist.url, + Self::Git(dist) => dist.url, + Self::Path(dist) => dist.url, + } + } +} + impl std::fmt::Display for SourceUrl<'_> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { diff --git a/crates/uv-requirements/src/lookahead.rs b/crates/uv-requirements/src/lookahead.rs index 1275352da5e4..c38704d75276 100644 --- a/crates/uv-requirements/src/lookahead.rs +++ b/crates/uv-requirements/src/lookahead.rs @@ -135,7 +135,8 @@ impl<'a, Context: BuildContext + Send + Sync> LookaheadResolver<'a, Context> { // Fetch the metadata for the distribution. let requires_dist = { // If the metadata is already in the index, return it. - if let Some(metadata) = self.index.get_metadata(&dist.package_id()) { + let id = dist.package_id(); + if let Some(metadata) = self.index.get_metadata(&id) { metadata.requires_dist.clone() } else { // Run the PEP 517 build process to extract metadata from the source distribution. @@ -149,8 +150,7 @@ impl<'a, Context: BuildContext + Send + Sync> LookaheadResolver<'a, Context> { })?; // Insert the metadata into the index. - self.index - .insert_metadata(dist.package_id(), metadata.clone()); + self.index.insert_metadata(id, metadata.clone()); // Insert the redirect into the index. if let Some(precise) = precise { diff --git a/crates/uv-requirements/src/unnamed.rs b/crates/uv-requirements/src/unnamed.rs index f4d237b2d461..6a8cdb919018 100644 --- a/crates/uv-requirements/src/unnamed.rs +++ b/crates/uv-requirements/src/unnamed.rs @@ -1,7 +1,6 @@ use std::borrow::Cow; use std::path::Path; use std::str::FromStr; -use std::sync::Arc; use anyhow::{Context, Result}; use configparser::ini::Ini; @@ -9,64 +8,69 @@ use futures::{StreamExt, TryStreamExt}; use serde::Deserialize; use tracing::debug; +use cache_key::CanonicalUrl; use distribution_filename::{SourceDistFilename, WheelFilename}; use distribution_types::{ - BuildableSource, DirectSourceUrl, GitSourceUrl, PathSourceUrl, RemoteSource, SourceUrl, + BuildableSource, DirectSourceUrl, GitSourceUrl, PackageId, PathSourceUrl, RemoteSource, + SourceUrl, }; use pep508_rs::{ Requirement, RequirementsTxtRequirement, Scheme, UnnamedRequirement, VersionOrUrl, }; use pypi_types::Metadata10; use uv_client::RegistryClient; -use uv_distribution::{Reporter, SourceDistributionBuilder}; +use uv_distribution::{DistributionDatabase, Reporter}; use uv_normalize::PackageName; +use uv_resolver::InMemoryIndex; use uv_types::BuildContext; /// Like [`RequirementsSpecification`], but with concrete names for all requirements. -pub struct NamedRequirementsResolver { +pub struct NamedRequirementsResolver<'a, Context: BuildContext + Send + Sync> { /// The requirements for the project. requirements: Vec, - /// The reporter to use when building source distributions. - reporter: Option>, + /// The in-memory index for resolving dependencies. + index: &'a InMemoryIndex, + /// The database for fetching and building distributions. + database: DistributionDatabase<'a, Context>, } -impl NamedRequirementsResolver { +impl<'a, Context: BuildContext + Send + Sync> NamedRequirementsResolver<'a, Context> { /// Instantiate a new [`NamedRequirementsResolver`] for a given set of requirements. - pub fn new(requirements: Vec) -> Self { + pub fn new( + requirements: Vec, + context: &'a Context, + client: &'a RegistryClient, + index: &'a InMemoryIndex, + ) -> Self { Self { requirements, - reporter: None, + index, + database: DistributionDatabase::new(client, context), } } /// Set the [`Reporter`] to use for this resolver. #[must_use] pub fn with_reporter(self, reporter: impl Reporter + 'static) -> Self { - let reporter: Arc = Arc::new(reporter); Self { - reporter: Some(reporter), + database: self.database.with_reporter(reporter), ..self } } /// Resolve any unnamed requirements in the specification. - pub async fn resolve( - self, - context: &T, - client: &RegistryClient, - ) -> Result> { - futures::stream::iter(self.requirements) + pub async fn resolve(self) -> Result> { + let Self { + requirements, + index, + database, + } = self; + futures::stream::iter(requirements) .map(|requirement| async { match requirement { RequirementsTxtRequirement::Pep508(requirement) => Ok(requirement), RequirementsTxtRequirement::Unnamed(requirement) => { - Self::resolve_requirement( - requirement, - context, - client, - self.reporter.clone(), - ) - .await + Self::resolve_requirement(requirement, index, &database).await } } }) @@ -76,11 +80,10 @@ impl NamedRequirementsResolver { } /// Infer the package name for a given "unnamed" requirement. - async fn resolve_requirement( + async fn resolve_requirement( requirement: UnnamedRequirement, - context: &T, - client: &RegistryClient, - reporter: Option>, + index: &InMemoryIndex, + database: &DistributionDatabase<'a, Context>, ) -> Result { // If the requirement is a wheel, extract the package name from the wheel filename. // @@ -231,20 +234,35 @@ impl NamedRequirementsResolver { } }; - // Run the PEP 517 build process to extract metadata from the source distribution. - let builder = if let Some(reporter) = reporter { - SourceDistributionBuilder::new(client, context).with_reporter(reporter) - } else { - SourceDistributionBuilder::new(client, context) - }; + // Fetch the metadata for the distribution. + let name = { + // If the metadata is already in the index, return it. + let id = PackageId::from_url(source.url()); + if let Some(metadata) = index.get_metadata(&id) { + metadata.name.clone() + } else { + // Run the PEP 517 build process to extract metadata from the source distribution. + let (metadata, precise) = database + .build_wheel_metadata(&BuildableSource::Url(source)) + .await + .context("Failed to build source distribution")?; - let metadata = builder - .download_and_build_metadata(&BuildableSource::Url(source)) - .await - .context("Failed to build source distribution")?; + let name = metadata.name.clone(); + + // Insert the metadata into the index. + index.insert_metadata(id, metadata); + + // Insert the redirect into the index. + if let Some(precise) = precise { + index.insert_redirect(CanonicalUrl::new(&requirement.url), precise); + } + + name + } + }; Ok(Requirement { - name: metadata.name, + name, extras: requirement.extras, version_or_url: Some(VersionOrUrl::Url(requirement.url)), marker: requirement.marker, diff --git a/crates/uv/src/commands/pip_compile.rs b/crates/uv/src/commands/pip_compile.rs index 4c2b793c5598..101d8138e2fc 100644 --- a/crates/uv/src/commands/pip_compile.rs +++ b/crates/uv/src/commands/pip_compile.rs @@ -263,10 +263,15 @@ pub(crate) async fn pip_compile( // Resolve the requirements from the provided sources. let requirements = { // Convert from unnamed to named requirements. - let mut requirements = NamedRequirementsResolver::new(requirements) - .with_reporter(ResolverReporter::from(printer)) - .resolve(&build_dispatch, &client) - .await?; + let mut requirements = NamedRequirementsResolver::new( + requirements, + &build_dispatch, + &client, + &top_level_index, + ) + .with_reporter(ResolverReporter::from(printer)) + .resolve() + .await?; // Resolve any source trees into requirements. if !source_trees.is_empty() { diff --git a/crates/uv/src/commands/pip_install.rs b/crates/uv/src/commands/pip_install.rs index 3479a3bc008d..461c5abb42f1 100644 --- a/crates/uv/src/commands/pip_install.rs +++ b/crates/uv/src/commands/pip_install.rs @@ -244,10 +244,11 @@ pub(crate) async fn pip_install( // Resolve the requirements from the provided sources. let requirements = { // Convert from unnamed to named requirements. - let mut requirements = NamedRequirementsResolver::new(requirements) - .with_reporter(ResolverReporter::from(printer)) - .resolve(&resolve_dispatch, &client) - .await?; + let mut requirements = + NamedRequirementsResolver::new(requirements, &resolve_dispatch, &client, &index) + .with_reporter(ResolverReporter::from(printer)) + .resolve() + .await?; // Resolve any source trees into requirements. if !source_trees.is_empty() { diff --git a/crates/uv/src/commands/pip_sync.rs b/crates/uv/src/commands/pip_sync.rs index e994a213fc9c..aea04aa6f34f 100644 --- a/crates/uv/src/commands/pip_sync.rs +++ b/crates/uv/src/commands/pip_sync.rs @@ -194,10 +194,11 @@ pub(crate) async fn pip_sync( // Convert from unnamed to named requirements. let requirements = { // Convert from unnamed to named requirements. - let mut requirements = NamedRequirementsResolver::new(requirements) - .with_reporter(ResolverReporter::from(printer)) - .resolve(&build_dispatch, &client) - .await?; + let mut requirements = + NamedRequirementsResolver::new(requirements, &build_dispatch, &client, &index) + .with_reporter(ResolverReporter::from(printer)) + .resolve() + .await?; // Resolve any source trees into requirements. if !source_trees.is_empty() {