From bed4bdbece02c295e8b66b001d4695f52d6bc5e3 Mon Sep 17 00:00:00 2001 From: konstin Date: Thu, 20 Jun 2024 20:33:21 +0200 Subject: [PATCH] Flatten requirements eagerly in `get_dependencies` Consider these two cases: A: ``` werkzeug==2.0.0 werkzeug @ https://files.pythonhosted.org/packages/ff/1d/960bb4017c68674a1cb099534840f18d3def3ce44aed12b5ed8b78e0153e/Werkzeug-2.0.0-py3-none-any.whl ``` B: ```toml dependencies = [ "iniconfig == 1.1.1 ; python_version < '3.12'", "iniconfig @ git+https://github.com/pytest-dev/iniconfig@93f5930e668c0d1ddf4597e38dd0dea4e2665e7a ; python_version >= '3.12'", ] ``` In the first case, `werkzeug==2.0.0` should be overridden by the url. In the second case `iniconfig == 1.1.1` is in a different fork and must remain a registry distribution. That means the conversion from `Requirement` to `PubGrubPackage` is dependent on the other requirements of the package. We can either look into the other packages immediately, or we can move the forking before the conversion to `PubGrubDependencies` instead of after. Either version requires a flat list of `Requirement`s to use. This refactoring gives us this list. --- crates/uv-configuration/src/constraints.rs | 6 +- .../uv-resolver/src/pubgrub/dependencies.rs | 293 +++--------------- crates/uv-resolver/src/resolver/mod.rs | 216 +++++++++++-- 3 files changed, 231 insertions(+), 284 deletions(-) diff --git a/crates/uv-configuration/src/constraints.rs b/crates/uv-configuration/src/constraints.rs index 187431536685..ffba0476d2c2 100644 --- a/crates/uv-configuration/src/constraints.rs +++ b/crates/uv-configuration/src/constraints.rs @@ -16,7 +16,11 @@ impl Constraints { constraints .entry(requirement.name.clone()) .or_default() - .push(requirement); + .push(Requirement { + // We add and apply constraints independent of their extras. + extras: vec![], + ..requirement + }); } Self(constraints) } diff --git a/crates/uv-resolver/src/pubgrub/dependencies.rs b/crates/uv-resolver/src/pubgrub/dependencies.rs index 80976a6891fa..fced9adb5a05 100644 --- a/crates/uv-resolver/src/pubgrub/dependencies.rs +++ b/crates/uv-resolver/src/pubgrub/dependencies.rs @@ -1,26 +1,19 @@ -use std::collections::BTreeMap; - -use either::Either; use itertools::Itertools; use pubgrub::range::Range; -use rustc_hash::FxHashSet; -use tracing::{trace, warn}; +use tracing::warn; use distribution_types::Verbatim; use pep440_rs::Version; -use pep508_rs::{MarkerEnvironment, MarkerTree}; use pypi_types::{ ParsedArchiveUrl, ParsedDirectoryUrl, ParsedGitUrl, ParsedPathUrl, ParsedUrl, Requirement, RequirementSource, }; -use uv_configuration::{Constraints, Overrides}; use uv_git::GitResolver; -use uv_normalize::{ExtraName, GroupName, PackageName}; +use uv_normalize::{ExtraName, PackageName}; -use crate::pubgrub::specifier::PubGrubSpecifier; use crate::pubgrub::{PubGrubPackage, PubGrubPackageInner}; use crate::resolver::{Locals, Urls}; -use crate::ResolveError; +use crate::{PubGrubSpecifier, ResolveError}; #[derive(Debug)] pub struct PubGrubDependencies(Vec<(PubGrubPackage, Range)>); @@ -29,41 +22,51 @@ impl PubGrubDependencies { /// Generate a set of PubGrub dependencies from a set of requirements. #[allow(clippy::too_many_arguments)] pub(crate) fn from_requirements( - requirements: &[Requirement], - dev_dependencies: &BTreeMap>, - constraints: &Constraints, - overrides: &Overrides, + flattened_requirements: &[&Requirement], source_name: Option<&PackageName>, - source_extra: Option<&ExtraName>, - source_dev: Option<&GroupName>, urls: &Urls, locals: &Locals, git: &GitResolver, - env: Option<&MarkerEnvironment>, - requires_python: Option<&MarkerTree>, - fork_markers: &MarkerTree, ) -> Result { - let mut dependencies = Vec::default(); - let mut seen = FxHashSet::default(); - - add_requirements( - requirements, - dev_dependencies, - constraints, - overrides, - source_name, - source_extra, - source_dev, - urls, - locals, - git, - env, - requires_python, - fork_markers, - &mut dependencies, - &mut seen, - )?; + let mut dependencies = Vec::new(); + for requirement in flattened_requirements { + // Add the package, plus any extra variants. + for result in std::iter::once(PubGrubRequirement::from_requirement( + requirement, + None, + urls, + locals, + git, + )) + .chain(requirement.extras.clone().into_iter().map(|extra| { + PubGrubRequirement::from_requirement(requirement, Some(extra), urls, locals, git) + })) { + let PubGrubRequirement { package, version } = result?; + + match &*package { + PubGrubPackageInner::Package { name, .. } => { + // Detect self-dependencies. + if source_name.is_some_and(|source_name| source_name == name) { + warn!("{name} has a dependency on itself"); + continue; + } + dependencies.push((package.clone(), version.clone())); + } + PubGrubPackageInner::Marker { .. } => { + dependencies.push((package.clone(), version.clone())); + } + PubGrubPackageInner::Extra { name, .. } => { + debug_assert!( + !source_name.is_some_and(|source_name| source_name == name), + "extras not flattened for {name}" + ); + dependencies.push((package.clone(), version.clone())); + } + _ => {} + } + } + } Ok(Self(dependencies)) } @@ -78,183 +81,6 @@ impl PubGrubDependencies { } } -/// Add a set of requirements to a list of dependencies. -#[allow(clippy::too_many_arguments)] -fn add_requirements( - requirements: &[Requirement], - dev_dependencies: &BTreeMap>, - constraints: &Constraints, - overrides: &Overrides, - source_name: Option<&PackageName>, - source_extra: Option<&ExtraName>, - source_dev: Option<&GroupName>, - urls: &Urls, - locals: &Locals, - git: &GitResolver, - env: Option<&MarkerEnvironment>, - requires_python: Option<&MarkerTree>, - fork_markers: &MarkerTree, - dependencies: &mut Vec<(PubGrubPackage, Range)>, - seen: &mut FxHashSet, -) -> Result<(), ResolveError> { - // Iterate over all declared requirements. - for requirement in overrides.apply(if let Some(source_dev) = source_dev { - Either::Left(dev_dependencies.get(source_dev).into_iter().flatten()) - } else { - Either::Right(requirements.iter()) - }) { - // If the requirement would not be selected with any Python version - // supported by the root, skip it. - if !satisfies_requires_python(requires_python, requirement) { - trace!( - "skipping {requirement} because of Requires-Python {requires_python}", - // OK because this filter only applies when there is a present - // Requires-Python specifier. - requires_python = requires_python.unwrap() - ); - continue; - } - // If we're in universal mode, `fork_markers` might correspond to a - // non-trivial marker expression that provoked the resolver to fork. - // In that case, we should ignore any dependency that cannot possibly - // satisfy the markers that provoked the fork. - if !possible_to_satisfy_markers(fork_markers, requirement) { - trace!("skipping {requirement} because of context resolver markers {fork_markers}"); - continue; - } - // If the requirement isn't relevant for the current platform, skip it. - match source_extra { - Some(source_extra) => { - // Only include requirements that are relevant for the current extra. - if requirement.evaluate_markers(env, &[]) { - continue; - } - if !requirement.evaluate_markers(env, std::slice::from_ref(source_extra)) { - continue; - } - } - None => { - if !requirement.evaluate_markers(env, &[]) { - continue; - } - } - } - - // Add the package, plus any extra variants. - for result in std::iter::once(PubGrubRequirement::from_requirement( - requirement, - None, - urls, - locals, - git, - )) - .chain(requirement.extras.clone().into_iter().map(|extra| { - PubGrubRequirement::from_requirement(requirement, Some(extra), urls, locals, git) - })) { - let PubGrubRequirement { package, version } = result?; - - match &*package { - PubGrubPackageInner::Package { name, .. } => { - // Detect self-dependencies. - if source_name.is_some_and(|source_name| source_name == name) { - warn!("{name} has a dependency on itself"); - continue; - } - - dependencies.push((package.clone(), version.clone())); - } - PubGrubPackageInner::Marker { .. } => { - dependencies.push((package.clone(), version.clone())); - } - PubGrubPackageInner::Extra { name, extra, .. } => { - // Recursively add the dependencies of the current package (e.g., `black` depending on - // `black[colorama]`). - if source_name.is_some_and(|source_name| source_name == name) { - if seen.insert(extra.clone()) { - add_requirements( - requirements, - dev_dependencies, - constraints, - overrides, - source_name, - Some(extra), - None, - urls, - locals, - git, - env, - requires_python, - fork_markers, - dependencies, - seen, - )?; - } - } else { - dependencies.push((package.clone(), version.clone())); - } - } - _ => {} - } - - // If the requirement was constrained, add those constraints. - for constraint in constraints.get(&requirement.name).into_iter().flatten() { - // If the requirement would not be selected with any Python - // version supported by the root, skip it. - if !satisfies_requires_python(requires_python, constraint) { - trace!( - "skipping {requirement} because of Requires-Python {requires_python}", - // OK because this filter only applies when there is a present - // Requires-Python specifier. - requires_python = requires_python.unwrap() - ); - continue; - } - // If we're in universal mode, `fork_markers` might correspond - // to a non-trivial marker expression that provoked the - // resolver to fork. In that case, we should ignore any - // dependency that cannot possibly satisfy the markers that - // provoked the fork. - if !possible_to_satisfy_markers(fork_markers, constraint) { - trace!( - "skipping {requirement} because of context resolver markers {fork_markers}" - ); - continue; - } - // If the requirement isn't relevant for the current platform, skip it. - match source_extra { - Some(source_extra) => { - if !constraint.evaluate_markers(env, std::slice::from_ref(source_extra)) { - continue; - } - } - None => { - if !constraint.evaluate_markers(env, &[]) { - continue; - } - } - } - - // Add the package. - let PubGrubRequirement { package, version } = - PubGrubRequirement::from_constraint(constraint, urls, locals, git)?; - - // Ignore self-dependencies. - if let PubGrubPackageInner::Package { name, .. } = &*package { - // Detect self-dependencies. - if source_name.is_some_and(|source_name| source_name == name) { - warn!("{name} has a dependency on itself"); - continue; - } - } - - dependencies.push((package.clone(), version.clone())); - } - } - } - - Ok(()) -} - /// Convert a [`PubGrubDependencies`] to a [`DependencyConstraints`]. impl From for Vec<(PubGrubPackage, Range)> { fn from(dependencies: PubGrubDependencies) -> Self { @@ -455,39 +281,4 @@ impl PubGrubRequirement { } } } - - /// Convert a constraint to a PubGrub-compatible package and range. - pub(crate) fn from_constraint( - constraint: &Requirement, - urls: &Urls, - locals: &Locals, - git: &GitResolver, - ) -> Result { - Self::from_requirement(constraint, None, urls, locals, git) - } -} - -/// Returns true if and only if the given requirement's marker expression has a -/// possible true value given the `requires_python` specifier given. -/// -/// While this is always called, a `requires_python` is only non-None when in -/// universal resolution mode. In non-universal mode, `requires_python` is -/// `None` and this always returns `true`. -fn satisfies_requires_python( - requires_python: Option<&MarkerTree>, - requirement: &Requirement, -) -> bool { - let Some(requires_python) = requires_python else { - return true; - }; - possible_to_satisfy_markers(requires_python, requirement) -} - -/// Returns true if and only if the given requirement's marker expression has a -/// possible true value given the `markers` expression given. -fn possible_to_satisfy_markers(markers: &MarkerTree, requirement: &Requirement) -> bool { - let Some(marker) = requirement.marker.as_ref() else { - return true; - }; - !crate::marker::is_disjoint(markers, marker) } diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index 195154855f3c..8bd3fb61960c 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -1,13 +1,14 @@ //! Given a set of requirements, find a set of compatible packages. use std::borrow::Cow; -use std::collections::BTreeMap; +use std::collections::{BTreeMap, VecDeque}; use std::fmt::{Display, Formatter}; use std::ops::Bound; use std::sync::Arc; -use std::thread; +use std::{iter, thread}; use dashmap::DashMap; +use either::Either; use futures::{FutureExt, StreamExt, TryFutureExt}; use itertools::Itertools; use pubgrub::error::PubGrubError; @@ -938,27 +939,26 @@ impl ResolverState Result { let (dependencies, name) = match &**package { PubGrubPackageInner::Root(_) => { - // Add the root requirements. - let dependencies = PubGrubDependencies::from_requirements( + let no_dev_deps = BTreeMap::default(); + let requirements = self.flatten_requirements( &self.requirements, - &BTreeMap::default(), - &self.constraints, - &self.overrides, + &no_dev_deps, None, None, None, + markers, + ); + + let dependencies = PubGrubDependencies::from_requirements( + &requirements, + None, &self.urls, &self.locals, &self.git, - self.markers.as_ref(), - self.requires_python.as_ref(), - markers, )?; + (dependencies, None) } - - PubGrubPackageInner::Python(_) => return Ok(Dependencies::Available(Vec::default())), - PubGrubPackageInner::Package { name, extra, @@ -1072,43 +1072,45 @@ impl ResolverState return Ok(Dependencies::Available(Vec::default())), // Add a dependency on both the marker and base package. PubGrubPackageInner::Marker { name, marker, url } => { @@ -1195,6 +1197,7 @@ impl ResolverState ResolverState( + &'a self, + dependencies: &'a [Requirement], + dev_dependencies: &'a BTreeMap>, + extra: Option<&'a ExtraName>, + dev: Option<&'a GroupName>, + name: Option<&PackageName>, + markers: &'a MarkerTree, + ) -> Vec<&'a Requirement> { + // Start with the requirements for the current extra of the package (for an extra + // requirement) or the non-extra (regular) dependencies (if extra is None), plus + // the constraints for the current package. + let regular_and_dev_dependencies = if let Some(dev) = dev { + Either::Left(dev_dependencies.get(dev).into_iter().flatten()) + } else { + Either::Right(dependencies.iter()) + }; + let mut requirements: Vec<&Requirement> = self + .requirements_for_extra(regular_and_dev_dependencies, extra, markers) + .collect(); + + // Check if there are recursive self inclusions and we need to go into the expensive branch. + if !requirements + .iter() + .any(|req| name == Some(&req.name) && !req.extras.is_empty()) + { + return requirements; + } + + // Transitively process all extras that are recursively included, starting with the current + // extra. + let mut seen = FxHashSet::default(); + let mut queue: VecDeque<_> = requirements + .iter() + .filter(|req| name == Some(&req.name)) + .flat_map(|req| &req.extras) + .collect(); + while let Some(extra) = queue.pop_front() { + if !seen.insert(extra) { + continue; + } + // Add each transitively included extra. + queue.extend( + self.requirements_for_extra(dependencies, Some(extra), markers) + .filter(|req| name == Some(&req.name)) + .flat_map(|req| &req.extras), + ); + // Add the requirements for that extra + requirements.extend(self.requirements_for_extra(dependencies, Some(extra), markers)); + } + // Drop all the self-requirements now that we flattened them out. + requirements.retain(|req| name != Some(&req.name)); + + requirements + } + + /// The set of the regular and dev dependencies, filtered by Python version, + /// the markers of this fork and the requested extra. + fn requirements_for_extra<'a>( + &'a self, + dependencies: impl IntoIterator, + extra: Option<&'a ExtraName>, + markers: &'a MarkerTree, + ) -> impl Iterator { + self.overrides + .apply(dependencies) + .flat_map(|requirement| { + iter::once(requirement).chain( + // If the requirement was constrained, add those constraints. + self.constraints + .get(&requirement.name) + .into_iter() + .flatten(), + ) + }) + .filter(move |requirement| { + // If the requirement would not be selected with any Python version + // supported by the root, skip it. + if !satisfies_requires_python(self.requires_python.as_ref(), requirement) { + trace!( + "skipping {requirement} because of Requires-Python {requires_python}", + // OK because this filter only applies when there is a present + // Requires-Python specifier. + requires_python = self.requires_python.as_ref().unwrap() + ); + return false; + } + + // If we're in universal mode, `fork_markers` might correspond to a + // non-trivial marker expression that provoked the resolver to fork. + // In that case, we should ignore any dependency that cannot possibly + // satisfy the markers that provoked the fork. + if !possible_to_satisfy_markers(markers, requirement) { + trace!("skipping {requirement} because of context resolver markers {markers}"); + return false; + } + + // If the requirement isn't relevant for the current platform, skip it. + match extra { + Some(source_extra) => { + // Only include requirements that are relevant for the current extra. + if requirement.evaluate_markers(self.markers.as_ref(), &[]) { + return false; + } + if !requirement.evaluate_markers( + self.markers.as_ref(), + std::slice::from_ref(source_extra), + ) { + return false; + } + } + None => { + if !requirement.evaluate_markers(self.markers.as_ref(), &[]) { + return false; + } + } + } + true + }) + } + /// Fetch the metadata for a stream of packages and versions. async fn fetch( self: Arc, @@ -2290,3 +2417,28 @@ impl<'a> PossibleFork<'a> { false } } + +/// Returns true if and only if the given requirement's marker expression has a +/// possible true value given the `requires_python` specifier given. +/// +/// While this is always called, a `requires_python` is only non-None when in +/// universal resolution mode. In non-universal mode, `requires_python` is +/// `None` and this always returns `true`. +fn satisfies_requires_python( + requires_python: Option<&MarkerTree>, + requirement: &Requirement, +) -> bool { + let Some(requires_python) = requires_python else { + return true; + }; + possible_to_satisfy_markers(requires_python, requirement) +} + +/// Returns true if and only if the given requirement's marker expression has a +/// possible true value given the `markers` expression given. +fn possible_to_satisfy_markers(markers: &MarkerTree, requirement: &Requirement) -> bool { + let Some(marker) = requirement.marker.as_ref() else { + return true; + }; + !crate::marker::is_disjoint(markers, marker) +}