Skip to content

Commit

Permalink
Avoid reusing cached downloaded binaries with --no-binary (astral-s…
Browse files Browse the repository at this point in the history
…h#7772)

## Summary

Historically, we've allowed the use of wheels that were downloaded from
PyPI even when the user passes `--no-binary`, if the wheel exists in the
cache. This PR modifies the cache lookup code such that we respect
`--no-build` and `--no-binary` in those paths.

Closes astral-sh#2154.
  • Loading branch information
charliermarsh authored Sep 29, 2024
1 parent c415251 commit 66d7ec5
Show file tree
Hide file tree
Showing 4 changed files with 201 additions and 63 deletions.
128 changes: 69 additions & 59 deletions crates/uv-distribution/src/index/registry_wheel_index.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,8 @@
use std::collections::hash_map::Entry;
use std::collections::BTreeMap;

use rustc_hash::FxHashMap;

use distribution_types::{CachedRegistryDist, Hashed, IndexLocations, IndexUrl};
use pep440_rs::Version;
use platform_tags::Tags;
use uv_cache::{Cache, CacheBucket, WheelCache};
use uv_fs::{directories, files, symlinks};
Expand All @@ -14,14 +12,23 @@ use uv_types::HashStrategy;
use crate::index::cached_wheel::CachedWheel;
use crate::source::{HttpRevisionPointer, LocalRevisionPointer, HTTP_REVISION, LOCAL_REVISION};

/// An entry in the [`RegistryWheelIndex`].
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct IndexEntry {
/// The cached distribution.
pub dist: CachedRegistryDist,
/// Whether the wheel was built from source (true), or downloaded from the registry directly (false).
pub built: bool,
}

/// A local index of distributions that originate from a registry, like `PyPI`.
#[derive(Debug)]
pub struct RegistryWheelIndex<'a> {
cache: &'a Cache,
tags: &'a Tags,
index_locations: &'a IndexLocations,
hasher: &'a HashStrategy,
index: FxHashMap<&'a PackageName, BTreeMap<Version, CachedRegistryDist>>,
index: FxHashMap<&'a PackageName, Vec<IndexEntry>>,
}

impl<'a> RegistryWheelIndex<'a> {
Expand All @@ -44,26 +51,12 @@ impl<'a> RegistryWheelIndex<'a> {
/// Return an iterator over available wheels for a given package.
///
/// If the package is not yet indexed, this will index the package by reading from the cache.
pub fn get(
&mut self,
name: &'a PackageName,
) -> impl Iterator<Item = (&Version, &CachedRegistryDist)> {
pub fn get(&mut self, name: &'a PackageName) -> impl Iterator<Item = &IndexEntry> {
self.get_impl(name).iter().rev()
}

/// Get the best wheel for the given package name and version.
///
/// If the package is not yet indexed, this will index the package by reading from the cache.
pub fn get_version(
&mut self,
name: &'a PackageName,
version: &Version,
) -> Option<&CachedRegistryDist> {
self.get_impl(name).get(version)
}

/// Get an entry in the index.
fn get_impl(&mut self, name: &'a PackageName) -> &BTreeMap<Version, CachedRegistryDist> {
fn get_impl(&mut self, name: &'a PackageName) -> &[IndexEntry] {
let versions = match self.index.entry(name) {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => entry.insert(Self::index(
Expand All @@ -84,8 +77,8 @@ impl<'a> RegistryWheelIndex<'a> {
tags: &Tags,
index_locations: &IndexLocations,
hasher: &HashStrategy,
) -> BTreeMap<Version, CachedRegistryDist> {
let mut versions = BTreeMap::new();
) -> Vec<IndexEntry> {
let mut entries = vec![];

// Collect into owned `IndexUrl`.
let flat_index_urls: Vec<IndexUrl> = index_locations
Expand Down Expand Up @@ -113,12 +106,19 @@ impl<'a> RegistryWheelIndex<'a> {
if let Some(wheel) =
CachedWheel::from_http_pointer(wheel_dir.join(file), cache)
{
// Enforce hash-checking based on the built distribution.
if wheel.satisfies(
hasher
.get_package(&wheel.filename.name, &wheel.filename.version),
) {
Self::add_wheel(wheel, tags, &mut versions);
if wheel.filename.compatibility(tags).is_compatible() {
// Enforce hash-checking based on the built distribution.
if wheel.satisfies(
hasher.get_package(
&wheel.filename.name,
&wheel.filename.version,
),
) {
entries.push(IndexEntry {
dist: wheel.into_registry_dist(),
built: false,
});
}
}
}
}
Expand All @@ -132,12 +132,19 @@ impl<'a> RegistryWheelIndex<'a> {
if let Some(wheel) =
CachedWheel::from_local_pointer(wheel_dir.join(file), cache)
{
// Enforce hash-checking based on the built distribution.
if wheel.satisfies(
hasher
.get_package(&wheel.filename.name, &wheel.filename.version),
) {
Self::add_wheel(wheel, tags, &mut versions);
if wheel.filename.compatibility(tags).is_compatible() {
// Enforce hash-checking based on the built distribution.
if wheel.satisfies(
hasher.get_package(
&wheel.filename.name,
&wheel.filename.version,
),
) {
entries.push(IndexEntry {
dist: wheel.into_registry_dist(),
built: false,
});
}
}
}
}
Expand Down Expand Up @@ -182,38 +189,41 @@ impl<'a> RegistryWheelIndex<'a> {
if let Some(revision) = revision {
for wheel_dir in symlinks(cache_shard.join(revision.id())) {
if let Some(wheel) = CachedWheel::from_built_source(wheel_dir) {
// Enforce hash-checking based on the source distribution.
if revision.satisfies(
hasher.get_package(&wheel.filename.name, &wheel.filename.version),
) {
Self::add_wheel(wheel, tags, &mut versions);
if wheel.filename.compatibility(tags).is_compatible() {
// Enforce hash-checking based on the source distribution.
if revision.satisfies(
hasher
.get_package(&wheel.filename.name, &wheel.filename.version),
) {
entries.push(IndexEntry {
dist: wheel.into_registry_dist(),
built: true,
});
}
}
}
}
}
}
}

versions
}

/// Add the [`CachedWheel`] to the index.
fn add_wheel(
wheel: CachedWheel,
tags: &Tags,
versions: &mut BTreeMap<Version, CachedRegistryDist>,
) {
let dist_info = wheel.into_registry_dist();

// Pick the wheel with the highest priority
let compatibility = dist_info.filename.compatibility(tags);
if let Some(existing) = versions.get_mut(&dist_info.filename.version) {
// Override if we have better compatibility
if compatibility > existing.filename.compatibility(tags) {
*existing = dist_info;
}
} else if compatibility.is_compatible() {
versions.insert(dist_info.filename.version.clone(), dist_info);
}
// Sort the cached distributions by (1) version, (2) compatibility, and (3) build status.
// We want the highest versions, with the greatest compatibility, that were built from source.
// at the end of the list.
entries.sort_unstable_by(|a, b| {
a.dist
.filename
.version
.cmp(&b.dist.filename.version)
.then_with(|| {
a.dist
.filename
.compatibility(tags)
.cmp(&b.dist.filename.compatibility(tags))
.then_with(|| a.built.cmp(&b.built))
})
});

entries
}
}
17 changes: 14 additions & 3 deletions crates/uv-installer/src/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,7 @@ impl<'a> Planner<'a> {

// Check if installation of a binary version of the package should be allowed.
let no_binary = build_options.no_binary_package(&requirement.name);
let no_build = build_options.no_build_package(&requirement.name);

let installed_dists = site_packages.remove_packages(&requirement.name);

Expand Down Expand Up @@ -143,9 +144,19 @@ impl<'a> Planner<'a> {
// Identify any cached distributions that satisfy the requirement.
match &requirement.source {
RequirementSource::Registry { specifier, .. } => {
if let Some((_version, distribution)) = registry_index
.get(&requirement.name)
.find(|(version, _)| specifier.contains(version))
if let Some(distribution) =
registry_index.get(&requirement.name).find_map(|entry| {
if !specifier.contains(&entry.dist.filename.version) {
return None;
};
if entry.built && no_build {
return None;
}
if !entry.built && no_binary {
return None;
}
Some(&entry.dist)
})
{
debug!("Requirement already cached: {distribution}");
cached.push(CachedDist::Registry(distribution.clone()));
Expand Down
2 changes: 1 addition & 1 deletion crates/uv/tests/lock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7473,7 +7473,7 @@ fn lock_find_links_http_sdist() -> Result<()> {
----- stdout -----

----- stderr -----
Prepared 1 package in [TIME]
Prepared 2 packages in [TIME]
Installed 2 packages in [TIME]
+ packaging==23.2
+ project==0.1.0 (from file://[TEMP_DIR]/)
Expand Down
117 changes: 117 additions & 0 deletions crates/uv/tests/pip_install.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2010,6 +2010,66 @@ fn install_only_binary_all_and_no_binary_all() {
context.assert_command("import anyio").failure();
}

/// Binary dependencies in the cache should be reused when the user provides `--no-build`.
#[test]
fn install_no_binary_cache() {
let context = TestContext::new("3.12");

// Install a binary distribution.
uv_snapshot!(
context.pip_install().arg("idna"),
@r###"
success: true
exit_code: 0
----- stdout -----
----- stderr -----
Resolved 1 package in [TIME]
Prepared 1 package in [TIME]
Installed 1 package in [TIME]
+ idna==3.6
"###
);

// Re-create the virtual environment.
context.venv().assert().success();

// Re-install. The distribution should be installed from the cache.
uv_snapshot!(
context.pip_install().arg("idna"),
@r###"
success: true
exit_code: 0
----- stdout -----
----- stderr -----
Resolved 1 package in [TIME]
Installed 1 package in [TIME]
+ idna==3.6
"###
);

// Re-create the virtual environment.
context.venv().assert().success();

// Install with `--no-binary`. The distribution should be built from source, despite a binary
// distribution being available in the cache.
uv_snapshot!(
context.pip_install().arg("idna").arg("--no-binary").arg(":all:"),
@r###"
success: true
exit_code: 0
----- stdout -----
----- stderr -----
Resolved 1 package in [TIME]
Prepared 1 package in [TIME]
Installed 1 package in [TIME]
+ idna==3.6
"###
);
}

/// Respect `--only-binary` flags in `requirements.txt`
#[test]
fn only_binary_requirements_txt() {
Expand Down Expand Up @@ -2124,6 +2184,63 @@ fn only_binary_editable_setup_py() {
);
}

#[test]
fn cache_priority() {
let context = TestContext::new("3.12");

// Install a specific `idna` version.
uv_snapshot!(
context.pip_install().arg("idna==3.6"),
@r###"
success: true
exit_code: 0
----- stdout -----
----- stderr -----
Resolved 1 package in [TIME]
Prepared 1 package in [TIME]
Installed 1 package in [TIME]
+ idna==3.6
"###
);

// Install a lower `idna` version.
uv_snapshot!(
context.pip_install().arg("idna==3.0"),
@r###"
success: true
exit_code: 0
----- stdout -----
----- stderr -----
Resolved 1 package in [TIME]
Prepared 1 package in [TIME]
Uninstalled 1 package in [TIME]
Installed 1 package in [TIME]
- idna==3.6
+ idna==3.0
"###
);

// Re-create the virtual environment.
context.venv().assert().success();

// Install `idna` without a version specifier.
uv_snapshot!(
context.pip_install().arg("idna"),
@r###"
success: true
exit_code: 0
----- stdout -----
----- stderr -----
Resolved 1 package in [TIME]
Installed 1 package in [TIME]
+ idna==3.6
"###
);
}

/// Install a package into a virtual environment, and ensuring that the executable permissions
/// are retained.
///
Expand Down

0 comments on commit 66d7ec5

Please sign in to comment.