Skip to content

Commit

Permalink
Respect requires-python when prefetching (#4900)
Browse files Browse the repository at this point in the history
## Summary

This is fallout from #4705. We need
to respect `requires-python` in the prefetch code to avoid building
unsupported distributions.

Closes #4898.
  • Loading branch information
charliermarsh authored Jul 8, 2024
1 parent 71c6a9f commit ac3a085
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 8 deletions.
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

0 comments on commit ac3a085

Please sign in to comment.