Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Respect requires-python when prefetching #4900

Merged
merged 1 commit into from
Jul 8, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 40 additions & 4 deletions crates/uv-resolver/src/resolver/batch_prefetch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,13 @@ use rustc_hash::FxHashMap;
use tokio::sync::mpsc::Sender;
use tracing::{debug, trace};

use distribution_types::DistributionMetadata;
use distribution_types::{CompatibleDist, DistributionMetadata};
use pep440_rs::Version;

use crate::candidate_selector::{CandidateDist, CandidateSelector};
use crate::candidate_selector::CandidateSelector;
use crate::pubgrub::{PubGrubPackage, PubGrubPackageInner};
use crate::resolver::Request;
use crate::{InMemoryIndex, ResolveError, VersionsResponse};
use crate::{InMemoryIndex, PythonRequirement, ResolveError, VersionsResponse};

enum BatchPrefetchStrategy {
/// Go through the next versions assuming the existing selection and its constraints
Expand Down Expand Up @@ -49,6 +49,7 @@ impl BatchPrefetcher {
next: &PubGrubPackage,
version: &Version,
current_range: &Range<Version>,
python_requirement: &PythonRequirement,
request_sink: &Sender<Request>,
index: &InMemoryIndex,
selector: &CandidateSelector,
Expand Down Expand Up @@ -128,14 +129,49 @@ impl BatchPrefetcher {
}
};

let CandidateDist::Compatible(dist) = candidate.dist() else {
let Some(dist) = candidate.compatible() else {
continue;
};

// Avoid prefetching source distributions, which could be expensive.
if !dist.prefetchable() {
continue;
}

// Avoid prefetching for distributions that don't satisfy the Python requirement.
match dist {
CompatibleDist::InstalledDist(_) => {}
CompatibleDist::SourceDist { sdist, .. }
| CompatibleDist::IncompatibleWheel { sdist, .. } => {
// Source distributions must meet both the _target_ Python version and the
// _installed_ Python version (to build successfully).
if let Some(requires_python) = sdist.file.requires_python.as_ref() {
if let Some(target) = python_requirement.target() {
if !target.is_compatible_with(requires_python) {
continue;
}
}
if !requires_python.contains(python_requirement.installed()) {
continue;
}
}
}
CompatibleDist::CompatibleWheel { wheel, .. } => {
// Wheels must meet the _target_ Python version.
if let Some(requires_python) = wheel.file.requires_python.as_ref() {
if let Some(target) = python_requirement.target() {
if !target.is_compatible_with(requires_python) {
continue;
}
} else {
if !requires_python.contains(python_requirement.installed()) {
continue;
}
}
}
}
};

let dist = dist.for_resolution();

// Emit a request to fetch the metadata for this version.
Expand Down
48 changes: 44 additions & 4 deletions crates/uv-resolver/src/resolver/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -364,6 +364,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
Self::pre_visit(
state.pubgrub.partial_solution.prioritized_packages(),
&self.urls,
&state.python_requirement,
&request_sink,
)?;
}
Expand Down Expand Up @@ -487,6 +488,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
&state.next,
&version,
term_intersection.unwrap_positive(),
&state.python_requirement,
&request_sink,
&self.index,
&self.selector,
Expand Down Expand Up @@ -721,6 +723,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
fn pre_visit<'data>(
packages: impl Iterator<Item = (&'data PubGrubPackage, &'data Range<Version>)>,
urls: &Urls,
python_requirement: &PythonRequirement,
request_sink: &Sender<Request>,
) -> Result<(), ResolveError> {
// Iterate over the potential packages, and fetch file metadata for any of them. These
Expand All @@ -740,7 +743,11 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
if urls.any_url(name) {
continue;
}
request_sink.blocking_send(Request::Prefetch(name.clone(), range.clone()))?;
request_sink.blocking_send(Request::Prefetch(
name.clone(),
range.clone(),
python_requirement.clone(),
))?;
}
Ok(())
}
Expand Down Expand Up @@ -1656,7 +1663,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
}

// Pre-fetch the package and distribution metadata.
Request::Prefetch(package_name, range) => {
Request::Prefetch(package_name, range, python_requirement) => {
// Wait for the package metadata to become available.
let versions_response = self
.index
Expand Down Expand Up @@ -1720,6 +1727,39 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
}
}

match dist {
CompatibleDist::InstalledDist(_) => {}
CompatibleDist::SourceDist { sdist, .. }
| CompatibleDist::IncompatibleWheel { sdist, .. } => {
// Source distributions must meet both the _target_ Python version and the
// _installed_ Python version (to build successfully).
if let Some(requires_python) = sdist.file.requires_python.as_ref() {
if let Some(target) = python_requirement.target() {
if !target.is_compatible_with(requires_python) {
return Ok(None);
}
}
if !requires_python.contains(python_requirement.installed()) {
return Ok(None);
}
}
}
CompatibleDist::CompatibleWheel { wheel, .. } => {
// Wheels must meet the _target_ Python version.
if let Some(requires_python) = wheel.file.requires_python.as_ref() {
if let Some(target) = python_requirement.target() {
if !target.is_compatible_with(requires_python) {
return Ok(None);
}
} else {
if !requires_python.contains(python_requirement.installed()) {
return Ok(None);
}
}
}
}
};

// Emit a request to fetch the metadata for this version.
if self.index.distributions().register(candidate.version_id()) {
let dist = dist.for_resolution().to_owned();
Expand Down Expand Up @@ -2213,7 +2253,7 @@ pub(crate) enum Request {
/// A request to fetch the metadata from an already-installed distribution.
Installed(InstalledDist),
/// A request to pre-fetch the metadata for a package and the best-guess distribution.
Prefetch(PackageName, Range<Version>),
Prefetch(PackageName, Range<Version>, PythonRequirement),
}

impl<'a> From<ResolvedDistRef<'a>> for Request {
Expand Down Expand Up @@ -2265,7 +2305,7 @@ impl Display for Request {
Self::Installed(dist) => {
write!(f, "Installed metadata {dist}")
}
Self::Prefetch(package_name, range) => {
Self::Prefetch(package_name, range, _) => {
write!(f, "Prefetch {package_name} {range}")
}
}
Expand Down
29 changes: 29 additions & 0 deletions crates/uv/tests/pip_compile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4544,6 +4544,35 @@ coverage = ["example[test]", "extras>=0.0.1,<=0.0.2"]
Ok(())
}

/// Respect `requires-python` when prefetching.
///
/// `voluptuous==0.15.1` requires Python 3.9 or later, so we should resolve to an earlier version
/// and avoiding building 0.15.1 at all.
#[test]
fn requires_python_prefetch() -> Result<()> {
let context = TestContext::new("3.8");
let requirements_in = context.temp_dir.child("requirements.in");
requirements_in.write_str("voluptuous<=0.15.1")?;

uv_snapshot!(context
.pip_compile()
.env_remove("UV_EXCLUDE_NEWER")
.arg("requirements.in"), @r###"
success: true
exit_code: 0
----- stdout -----
# This file was autogenerated by uv via the following command:
# uv pip compile --cache-dir [CACHE_DIR] requirements.in
voluptuous==0.14.2
# via -r requirements.in
----- stderr -----
Resolved 1 package in [TIME]
"###);

Ok(())
}

/// Use an existing resolution for `black==23.10.1`, with stale versions of `click` and `pathspec`.
/// Nothing should change.
#[test]
Expand Down
Loading