Track file index (#452)

Track the index (or at least its url) where we got a file from across
the source code.

Fixes #448
This commit is contained in:
konsti 2023-11-20 09:48:16 +01:00 committed by GitHub
parent 6fd582f8b9
commit 46bb18f06e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 86 additions and 39 deletions

View file

@ -17,7 +17,7 @@ use platform_tags::{TagPriority, Tags};
use puffin_client::RegistryClient;
use puffin_interpreter::InterpreterInfo;
use puffin_normalize::PackageName;
use pypi_types::{File, SimpleJson};
use pypi_types::{File, IndexUrl, SimpleJson};
use crate::error::ResolveError;
use crate::resolution::Resolution;
@ -68,7 +68,9 @@ impl<'a> DistFinder<'a> {
Request::Package(requirement) => self
.client
.simple(requirement.name.clone())
.map_ok(move |metadata| Response::Package(requirement, metadata)),
.map_ok(move |(index, metadata)| {
Response::Package(requirement, index, metadata)
}),
})
.buffer_unordered(32)
.ready_chunks(32);
@ -104,9 +106,10 @@ impl<'a> DistFinder<'a> {
for result in chunk {
let result: Response = result?;
match result {
Response::Package(requirement, metadata) => {
Response::Package(requirement, index, metadata) => {
// Pick a version that satisfies the requirement.
let Some(distribution) = self.select(&requirement, metadata.files) else {
let Some(distribution) = self.select(&requirement, &index, metadata.files)
else {
return Err(ResolveError::NotFound(requirement));
};
@ -134,7 +137,12 @@ impl<'a> DistFinder<'a> {
}
/// select a version that satisfies the requirement, preferring wheels to source distributions.
fn select(&self, requirement: &Requirement, files: Vec<File>) -> Option<Dist> {
fn select(
&self,
requirement: &Requirement,
index: &IndexUrl,
files: Vec<File>,
) -> Option<Dist> {
let mut best_version: Option<Version> = None;
let mut best_wheel: Option<(Dist, TagPriority)> = None;
let mut best_sdist: Option<Dist> = None;
@ -173,7 +181,7 @@ impl<'a> DistFinder<'a> {
.map_or(true, |(.., existing)| priority > *existing)
{
best_wheel = Some((
Dist::from_registry(wheel.name, wheel.version, file),
Dist::from_registry(wheel.name, wheel.version, file, index.clone()),
priority,
));
}
@ -197,7 +205,12 @@ impl<'a> DistFinder<'a> {
if requirement.is_satisfied_by(&sdist.version) {
best_version = Some(sdist.version.clone());
best_sdist = Some(Dist::from_registry(sdist.name, sdist.version, file));
best_sdist = Some(Dist::from_registry(
sdist.name,
sdist.version,
file,
index.clone(),
));
}
}
}
@ -216,7 +229,7 @@ enum Request {
#[derive(Debug)]
enum Response {
/// The returned metadata for a package.
Package(Requirement, SimpleJson),
Package(Requirement, IndexUrl, SimpleJson),
}
pub trait Reporter: Send + Sync {

View file

@ -14,7 +14,7 @@ use distribution_types::{BuiltDist, Dist, Metadata, SourceDist};
use pep440_rs::{Version, VersionSpecifier, VersionSpecifiers};
use pep508_rs::{Requirement, VersionOrUrl};
use puffin_normalize::PackageName;
use pypi_types::File;
use pypi_types::{File, IndexUrl};
use crate::pubgrub::{PubGrubPackage, PubGrubPriority, PubGrubVersion};
@ -58,7 +58,7 @@ impl Graph {
/// Create a new graph from the resolved `PubGrub` state.
pub fn from_state(
selection: &SelectedDependencies<PubGrubPackage, PubGrubVersion>,
pins: &FxHashMap<PackageName, FxHashMap<Version, File>>,
pins: &FxHashMap<PackageName, FxHashMap<Version, (IndexUrl, File)>>,
redirects: &WaitMap<Url, Url>,
state: &State<PubGrubPackage, Range<PubGrubVersion>, PubGrubPriority>,
) -> Self {
@ -73,12 +73,13 @@ impl Graph {
match package {
PubGrubPackage::Package(package_name, None, None) => {
let version = Version::from(version.clone());
let file = pins
let (index, file) = pins
.get(package_name)
.and_then(|versions| versions.get(&version))
.unwrap()
.clone();
let pinned_package = Dist::from_registry(package_name.clone(), version, file);
let pinned_package =
Dist::from_registry(package_name.clone(), version, file, index);
let index = graph.add_node(pinned_package);
inverse.insert(package_name, index);

View file

@ -27,7 +27,7 @@ use puffin_cache::CanonicalUrl;
use puffin_client::RegistryClient;
use puffin_normalize::{ExtraName, PackageName};
use puffin_traits::BuildContext;
use pypi_types::{File, Metadata21, SimpleJson};
use pypi_types::{File, IndexUrl, Metadata21, SimpleJson};
use crate::candidate_selector::CandidateSelector;
use crate::distribution::{BuiltDistFetcher, SourceDistFetcher, SourceDistributionReporter};
@ -330,7 +330,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
let Some(entry) = self.index.packages.get(package_name) else {
continue;
};
let version_map = entry.value();
let (index, version_map) = entry.value();
// Try to find a compatible version. If there aren't any compatible versions,
// short-circuit and return `None`.
@ -345,6 +345,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
candidate.package_name,
candidate.version.into(),
candidate.file.into(),
index.clone(),
);
request_sink.unbounded_send(Request::Dist(distribution))?;
}
@ -358,7 +359,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
&self,
package: &PubGrubPackage,
range: &Range<PubGrubVersion>,
pins: &mut FxHashMap<PackageName, FxHashMap<pep440_rs::Version, File>>,
pins: &mut FxHashMap<PackageName, FxHashMap<pep440_rs::Version, (IndexUrl, File)>>,
in_flight: &mut InFlight,
request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
) -> Result<Option<PubGrubVersion>, ResolveError> {
@ -405,7 +406,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
PubGrubPackage::Package(package_name, _extra, None) => {
// Wait for the metadata to be available.
let entry = self.index.packages.wait(package_name).await.unwrap();
let version_map = entry.value();
let (index, version_map) = entry.value();
debug!("Searching for a compatible version of {package_name} ({range})");
@ -428,7 +429,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
.or_default()
.insert(
candidate.version.clone().into(),
candidate.file.clone().into(),
(index.clone(), candidate.file.clone().into()),
);
let version = candidate.version.clone();
@ -439,6 +440,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
candidate.package_name,
candidate.version.into(),
candidate.file.into(),
index.clone(),
);
request_sink.unbounded_send(Request::Dist(distribution))?;
}
@ -453,7 +455,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
&self,
package: &PubGrubPackage,
version: &PubGrubVersion,
pins: &mut FxHashMap<PackageName, FxHashMap<pep440_rs::Version, File>>,
pins: &mut FxHashMap<PackageName, FxHashMap<pep440_rs::Version, (IndexUrl, File)>>,
priorities: &mut PubGrubPriorities,
in_flight: &mut InFlight,
request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
@ -498,7 +500,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
.unwrap(),
None => {
let versions = pins.get(package_name).unwrap();
let file = versions.get(version.into()).unwrap();
let (_index, file) = versions.get(version.into()).unwrap();
self.index
.distributions
.wait(&file.distribution_id())
@ -550,7 +552,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
while let Some(response) = response_stream.next().await {
match response? {
Response::Package(package_name, metadata) => {
Response::Package(package_name, index, metadata) => {
trace!("Received package metadata for: {package_name}");
let version_map = VersionMap::from_metadata(
metadata,
@ -559,7 +561,9 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
self.build_context.interpreter_info().version(),
self.exclude_newer.as_ref(),
);
self.index.packages.insert(package_name, version_map);
self.index
.packages
.insert(package_name, (index, version_map));
}
Response::Dist(Dist::Built(distribution), metadata, ..) => {
trace!("Received built distribution metadata for: {distribution}");
@ -596,7 +600,9 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
Request::Package(package_name) => {
self.client
.simple(package_name.clone())
.map_ok(move |metadata| Response::Package(package_name, metadata))
.map_ok(move |(index, metadata)| {
Response::Package(package_name, index, metadata)
})
.map_err(ResolveError::Client)
.await
}
@ -803,7 +809,7 @@ enum Request {
#[allow(clippy::large_enum_variant)]
enum Response {
/// The returned metadata for a package hosted on a registry.
Package(PackageName, SimpleJson),
Package(PackageName, IndexUrl, SimpleJson),
/// The returned metadata for a distribution.
Dist(Dist, Metadata21, Option<Url>),
}
@ -839,8 +845,9 @@ impl InFlight {
/// 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 package name to the metadata for that package and the index where the metadata
/// came from.
packages: WaitMap<PackageName, (IndexUrl, VersionMap)>,
/// A map from distribution SHA to metadata for that distribution.
distributions: WaitMap<String, Metadata21>,