Use a single requirements iterator in sync (#71)

This commit is contained in:
Charlie Marsh 2023-10-08 23:29:38 -04:00 committed by GitHub
parent ba72950546
commit 485b1dceb6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 90 additions and 88 deletions

View file

@ -52,7 +52,7 @@ pub(crate) async fn compile(src: &Path, cache: Option<&Path>) -> Result<ExitStat
// Resolve the dependencies.
let resolution = puffin_resolver::resolve(
&requirements,
requirements.iter(),
markers,
&tags,
&client,

View file

@ -8,7 +8,7 @@ use tracing::{debug, info};
use platform_host::Platform;
use platform_tags::Tags;
use puffin_client::PypiClientBuilder;
use puffin_installer::{Distribution, RemoteDistribution};
use puffin_installer::{Distribution, LocalDistribution, LocalIndex, RemoteDistribution};
use puffin_interpreter::{PythonExecutable, SitePackages};
use puffin_package::package_name::PackageName;
use puffin_package::requirements::Requirements;
@ -42,25 +42,53 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) ->
python.executable().display()
);
// Remove any already-installed packages.
let requirements = if flags.intersects(SyncFlags::IGNORE_INSTALLED) {
requirements
// Determine the current environment markers.
let markers = python.markers();
let tags = Tags::from_env(python.platform(), python.simple_version())?;
// Index all the already-installed packages in site-packages.
let site_packages = if flags.intersects(SyncFlags::IGNORE_INSTALLED) {
SitePackages::default()
} else {
let site_packages = SitePackages::from_executable(&python).await?;
requirements.filter(|requirement| {
let package = PackageName::normalize(&requirement.name);
if let Some(version) = site_packages.get(&package) {
#[allow(clippy::print_stdout)]
{
info!("Requirement already satisfied: {package} ({version})");
}
false
} else {
true
}
})
SitePackages::from_executable(&python).await?
};
// Index all the already-downloaded wheels in the cache.
let local_index = if let Some(cache) = cache {
LocalIndex::from_directory(cache).await?
} else {
LocalIndex::default()
};
let requirements = requirements
.iter()
.filter_map(|requirement| {
let package = PackageName::normalize(&requirement.name);
// Filter out already-installed packages.
if let Some(version) = site_packages.get(&package) {
info!("Requirement already satisfied: {package} ({version})");
return None;
}
// Identify any locally-available distributions that satisfy the requirement.
if let Some(distribution) = local_index
.get(&package)
.filter(|dist| requirement.is_satisfied_by(dist.version()))
{
debug!(
"Requirement already cached: {} ({})",
distribution.name(),
distribution.version()
);
return Some(Requirement::Local(distribution.clone()));
}
debug!("Identified uncached requirement: {}", requirement);
Some(Requirement::Remote(requirement.clone()))
})
.collect::<Vec<_>>();
if requirements.is_empty() {
let s = if initial_requirements == 1 { "" } else { "s" };
info!(
@ -72,42 +100,7 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) ->
return Ok(ExitStatus::Success);
}
// Detect any cached wheels.
let (uncached, cached) = if let Some(cache) = cache {
let mut cached = Vec::with_capacity(requirements.len());
let mut uncached = Vec::with_capacity(requirements.len());
let index = puffin_installer::LocalIndex::from_directory(cache).await?;
for requirement in requirements {
let package = PackageName::normalize(&requirement.name);
if let Some(distribution) = index
.get(&package)
.filter(|dist| requirement.is_satisfied_by(dist.version()))
{
debug!(
"Requirement already cached: {} ({})",
distribution.name(),
distribution.version()
);
cached.push(distribution.clone());
} else {
debug!("Identified uncached requirement: {}", requirement);
uncached.push(requirement);
}
}
(Requirements::new(uncached), cached)
} else {
(requirements, Vec::new())
};
// Determine the current environment markers.
let markers = python.markers();
// Determine the compatible platform tags.
let tags = Tags::from_env(python.platform(), python.simple_version())?;
// Instantiate a client.
// Resolve the dependencies.
let client = {
let mut pypi_client = PypiClientBuilder::default();
if let Some(cache) = cache {
@ -115,25 +108,27 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) ->
}
pypi_client.build()
};
let resolution = puffin_resolver::resolve(
requirements
.iter()
.filter_map(|requirement| match requirement {
Requirement::Remote(requirement) => Some(requirement),
Requirement::Local(_) => None,
}),
markers,
&tags,
&client,
puffin_resolver::ResolveFlags::NO_DEPS,
)
.await?;
// Resolve the dependencies.
let resolution = if uncached.is_empty() {
puffin_resolver::Resolution::empty()
} else {
puffin_resolver::resolve(
&uncached,
markers,
&tags,
&client,
puffin_resolver::ResolveFlags::NO_DEPS,
)
.await?
};
// Install into the current environment.
let wheels = cached
// Install the resolved distributions.
let wheels = requirements
.into_iter()
.map(|local| Ok(Distribution::Local(local)))
.filter_map(|requirement| match requirement {
Requirement::Remote(_) => None,
Requirement::Local(distribution) => Some(Ok(Distribution::Local(distribution))),
})
.chain(
resolution
.into_files()
@ -152,3 +147,11 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) ->
Ok(ExitStatus::Success)
}
#[derive(Debug)]
enum Requirement {
/// A requirement that must be downloaded from PyPI.
Remote(pep508_rs::Requirement),
/// A requirement that is already available locally.
Local(LocalDistribution),
}

View file

@ -1,21 +1,21 @@
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::path::Path;
use anyhow::Result;
use crate::cache::WheelCache;
use puffin_package::package_name::PackageName;
use crate::cache::WheelCache;
use crate::distribution::LocalDistribution;
/// A local index of cached distributions.
#[derive(Debug)]
pub struct LocalIndex(BTreeMap<PackageName, LocalDistribution>);
#[derive(Debug, Default)]
pub struct LocalIndex(HashMap<PackageName, LocalDistribution>);
impl LocalIndex {
/// Build an index of cached distributions from a directory.
pub async fn from_directory(path: &Path) -> Result<Self> {
let mut index = BTreeMap::new();
let mut index = HashMap::new();
let cache = WheelCache::new(path);
let Ok(mut dir) = cache.read_dir().await else {

View file

@ -1,4 +1,4 @@
pub use distribution::{Distribution, RemoteDistribution};
pub use distribution::{Distribution, LocalDistribution, RemoteDistribution};
pub use index::LocalIndex;
pub use install::install;

View file

@ -9,7 +9,7 @@ use puffin_package::package_name::PackageName;
use crate::PythonExecutable;
#[derive(Debug)]
#[derive(Debug, Default)]
pub struct SitePackages(BTreeMap<PackageName, Version>);
impl SitePackages {

View file

@ -14,17 +14,12 @@ use platform_tags::Tags;
use puffin_client::{File, PypiClient, SimpleJson};
use puffin_package::metadata::Metadata21;
use puffin_package::package_name::PackageName;
use puffin_package::requirements::Requirements;
use wheel_filename::WheelFilename;
#[derive(Debug)]
#[derive(Debug, Default)]
pub struct Resolution(HashMap<PackageName, PinnedPackage>);
impl Resolution {
pub fn empty() -> Self {
Self(HashMap::new())
}
/// Iterate over the pinned packages in this resolution.
pub fn iter(&self) -> impl Iterator<Item = (&PackageName, &PinnedPackage)> {
self.0.iter()
@ -86,7 +81,7 @@ impl<T> From<futures::channel::mpsc::TrySendError<T>> for ResolveError {
/// Resolve a set of requirements into a set of pinned versions.
pub async fn resolve(
requirements: &Requirements,
requirements: impl Iterator<Item = &Requirement>,
markers: &MarkerEnvironment,
tags: &Tags,
client: &PypiClient,
@ -116,16 +111,20 @@ pub async fn resolve(
.ready_chunks(32);
// Push all the requirements into the package sink.
let mut in_flight: HashSet<PackageName> = HashSet::with_capacity(requirements.len());
for requirement in requirements.iter() {
let mut in_flight: HashSet<PackageName> = HashSet::new();
for requirement in requirements {
debug!("--> adding root dependency: {}", requirement);
package_sink.unbounded_send(Request::Package(requirement.clone()))?;
in_flight.insert(PackageName::normalize(&requirement.name));
}
if in_flight.is_empty() {
return Ok(Resolution::default());
}
// Resolve the requirements.
let mut resolution: HashMap<PackageName, PinnedPackage> =
HashMap::with_capacity(requirements.len());
HashMap::with_capacity(in_flight.len());
while let Some(chunk) = package_stream.next().await {
for result in chunk {