From 7b6a5159f1aacfd8ffcb0bf3854b5df90c4f5dd2 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 16 Mar 2026 08:19:00 -0400 Subject: [PATCH] Make resolver dependency edges source-aware across forks (#18435) ## Summary This is a non-behavior-changing refactor that puts direct source information on individual `PubGrubDependency` edges. The resulting code is, in my opinion, a bit simpler with clearer abstractions and more consistent handling. --- .../uv-resolver/src/pubgrub/dependencies.rs | 117 +++++++++++++---- crates/uv-resolver/src/pubgrub/mod.rs | 2 +- crates/uv-resolver/src/resolver/fork_map.rs | 121 ++++++++++++++++-- crates/uv-resolver/src/resolver/indexes.rs | 34 +---- crates/uv-resolver/src/resolver/mod.rs | 23 ++-- crates/uv-resolver/src/resolver/system.rs | 4 +- crates/uv-resolver/src/resolver/urls.rs | 34 ++--- 7 files changed, 236 insertions(+), 99 deletions(-) diff --git a/crates/uv-resolver/src/pubgrub/dependencies.rs b/crates/uv-resolver/src/pubgrub/dependencies.rs index 7ef848ab2..82dc238ed 100644 --- a/crates/uv-resolver/src/pubgrub/dependencies.rs +++ b/crates/uv-resolver/src/pubgrub/dependencies.rs @@ -4,9 +4,10 @@ use std::iter; use either::Either; use pubgrub::Ranges; -use uv_distribution_types::{Requirement, RequirementSource}; +use uv_distribution_types::{IndexMetadata, Requirement, RequirementSource}; use uv_normalize::{ExtraName, GroupName, PackageName}; use uv_pep440::{Version, VersionSpecifiers}; +use uv_pep508::RequirementOrigin; use uv_pypi_types::{ ConflictItemRef, Conflicts, ParsedArchiveUrl, ParsedDirectoryUrl, ParsedGitUrl, ParsedPathUrl, ParsedUrl, VerbatimParsedUrl, @@ -14,6 +15,70 @@ use uv_pypi_types::{ use crate::pubgrub::{PubGrubPackage, PubGrubPackageInner}; +/// The source constraint carried by a single dependency edge. +/// +/// Most dependency edges are source-agnostic and use [`DependencySource::Unspecified`]. Direct +/// URLs and group-scoped explicit indexes use a concrete source so fork construction can keep +/// that source information attached to the edge that introduced it. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub(crate) enum DependencySource { + /// The dependency does not carry an edge-local source constraint. + #[default] + Unspecified, + /// The dependency was introduced by a direct URL-like requirement. + Url(Box), + /// The dependency was introduced by a requirement pinned to an explicit index. + ExplicitIndex(IndexMetadata), +} + +impl DependencySource { + /// Derive the edge-local source constraint from a requirement. + /// + /// Registry requirements only carry a source here when they are tied to a group-scoped + /// explicit index. Direct URL-like requirements always preserve their verbatim URL. + pub(crate) fn from_requirement(requirement: &Requirement) -> Self { + match &requirement.source { + RequirementSource::Registry { index, .. } + if matches!( + requirement.origin.as_ref(), + Some(RequirementOrigin::Group(_, Some(_), _)) + ) => + { + index + .clone() + .map(Self::ExplicitIndex) + .unwrap_or(Self::Unspecified) + } + RequirementSource::Registry { .. } => Self::Unspecified, + RequirementSource::Url { .. } + | RequirementSource::Git { .. } + | RequirementSource::Path { .. } + | RequirementSource::Directory { .. } => requirement + .source + .to_verbatim_parsed_url() + .map(Box::new) + .map(Self::Url) + .unwrap_or(Self::Unspecified), + } + } + + /// Return the direct URL attached to this source, if any. + pub(crate) fn verbatim_url(&self) -> Option<&VerbatimParsedUrl> { + match self { + Self::Url(url) => Some(url.as_ref()), + Self::Unspecified | Self::ExplicitIndex(_) => None, + } + } + + /// Return the explicit index attached to this source, if any. + pub(crate) fn explicit_index(&self) -> Option<&IndexMetadata> { + match self { + Self::ExplicitIndex(index) => Some(index), + Self::Unspecified | Self::Url(_) => None, + } + } +} + #[derive(Clone, Debug, Eq, PartialEq)] pub(crate) struct PubGrubDependency { pub(crate) package: PubGrubPackage, @@ -34,10 +99,12 @@ pub(crate) struct PubGrubDependency { /// we introduce this parent field to enable "delayed" filtering. pub(crate) parent: Option, - /// This field is set if the [`Requirement`] had a URL. We still use a URL from [`Urls`] - /// even if this field is None where there is an override with a URL or there is a different - /// requirement or constraint for the same package that has a URL. - pub(crate) url: Option, + /// The direct source constraint attached to this dependency edge. + /// + /// This is only populated when the edge itself needs source identity, e.g. for direct URLs + /// or group-scoped explicit indexes. Manifest-wide URL and index constraints are still applied + /// separately via `Urls` and `Indexes`. + pub(crate) source: DependencySource, } impl PubGrubDependency { @@ -104,7 +171,7 @@ impl PubGrubDependency { let PubGrubRequirement { package, version, - url, + source, } = pubgrub_requirement; match &*package { PubGrubPackageInner::Package { .. } => Self { @@ -115,7 +182,7 @@ impl PubGrubDependency { } else { None }, - url, + source, }, PubGrubPackageInner::Marker { .. } => Self { package, @@ -125,7 +192,7 @@ impl PubGrubDependency { } else { None }, - url, + source, }, PubGrubPackageInner::Extra { name, .. } => { if group_name.is_none() { @@ -138,7 +205,7 @@ impl PubGrubDependency { package, version, parent: None, - url, + source, } } PubGrubPackageInner::Group { name, .. } => { @@ -152,7 +219,7 @@ impl PubGrubDependency { package, version, parent: None, - url, + source, } } PubGrubPackageInner::Root(_) => unreachable!("Root package in dependencies"), @@ -178,10 +245,18 @@ impl PubGrubDependency { pub(crate) struct PubGrubRequirement { pub(crate) package: PubGrubPackage, pub(crate) version: Ranges, - pub(crate) url: Option, + pub(crate) source: DependencySource, } impl PubGrubRequirement { + fn package_for_requirement( + requirement: &Requirement, + extra: Option, + group: Option, + ) -> PubGrubPackage { + PubGrubPackage::from_package(requirement.name.clone(), extra, group, requirement.marker) + } + /// Convert a [`Requirement`] to a PubGrub-compatible package and range, while returning the URL /// on the [`Requirement`], if any. pub(crate) fn from_requirement( @@ -244,17 +319,12 @@ impl PubGrubRequirement { }; Self { - package: PubGrubPackage::from_package( - requirement.name.clone(), - extra, - group, - requirement.marker, - ), + package: Self::package_for_requirement(requirement, extra, group), version: Ranges::full(), - url: Some(VerbatimParsedUrl { + source: DependencySource::Url(Box::new(VerbatimParsedUrl { parsed_url, verbatim: verbatim_url.clone(), - }), + })), } } @@ -265,13 +335,8 @@ impl PubGrubRequirement { requirement: &Requirement, ) -> Self { Self { - package: PubGrubPackage::from_package( - requirement.name.clone(), - extra, - group, - requirement.marker, - ), - url: None, + package: Self::package_for_requirement(requirement, extra, group), + source: DependencySource::from_requirement(requirement), version: Ranges::from(specifier.clone()), } } diff --git a/crates/uv-resolver/src/pubgrub/mod.rs b/crates/uv-resolver/src/pubgrub/mod.rs index a0c96b150..fbc87433d 100644 --- a/crates/uv-resolver/src/pubgrub/mod.rs +++ b/crates/uv-resolver/src/pubgrub/mod.rs @@ -1,4 +1,4 @@ -pub(crate) use crate::pubgrub::dependencies::PubGrubDependency; +pub(crate) use crate::pubgrub::dependencies::{DependencySource, PubGrubDependency}; pub use crate::pubgrub::package::{PubGrubPackage, PubGrubPackageInner, PubGrubPython}; pub(crate) use crate::pubgrub::priority::{PubGrubPriorities, PubGrubPriority, PubGrubTiebreaker}; pub(crate) use crate::pubgrub::report::PubGrubReportFormatter; diff --git a/crates/uv-resolver/src/resolver/fork_map.rs b/crates/uv-resolver/src/resolver/fork_map.rs index 15f29856d..9a1d7bf53 100644 --- a/crates/uv-resolver/src/resolver/fork_map.rs +++ b/crates/uv-resolver/src/resolver/fork_map.rs @@ -1,10 +1,12 @@ use rustc_hash::FxHashMap; -use uv_distribution_types::Requirement; +use uv_distribution_types::{Requirement, RequirementSource}; use uv_normalize::PackageName; -use uv_pep508::MarkerTree; +use uv_pep508::{MarkerTree, RequirementOrigin}; +use uv_pypi_types::{ConflictItem, ConflictItemRef, ConflictKind}; use crate::ResolverEnvironment; +use crate::universal_marker::{ConflictMarker, UniversalMarker}; /// A set of package names associated with a given fork. pub(crate) type ForkSet = ForkMap<()>; @@ -17,7 +19,63 @@ pub(crate) struct ForkMap(FxHashMap>>); #[derive(Debug, Clone)] struct Entry { value: T, + scope: ForkScope, +} + +/// The fork visibility of an entry. +#[derive(Debug, Clone, Eq, PartialEq)] +struct ForkScope { marker: MarkerTree, + conflict: Option, +} + +impl ForkScope { + /// Derive the scope under which a requirement should be visible in forked resolution. + /// + /// Group conflicts are folded into the marker so group-scoped entries only appear in forks + /// where that group is active. + fn from_requirement(requirement: &Requirement) -> Self { + let conflict = Self::conflict_for_requirement(requirement); + let marker = conflict + .as_ref() + .filter(|conflict_item| matches!(conflict_item.kind(), ConflictKind::Group(_))) + .map_or(requirement.marker, |conflict_item| { + UniversalMarker::new( + requirement.marker.without_extras(), + ConflictMarker::from_conflict_item(conflict_item), + ) + .combined() + }); + Self { marker, conflict } + } + + fn conflict_for_requirement(requirement: &Requirement) -> Option { + let conflict = match &requirement.source { + RequirementSource::Registry { conflict, .. } => conflict.clone(), + RequirementSource::Url { .. } + | RequirementSource::Git { .. } + | RequirementSource::Path { .. } + | RequirementSource::Directory { .. } => None, + }; + conflict.or_else(|| match requirement.origin.as_ref() { + Some(RequirementOrigin::Group(_, Some(project_name), group)) => { + Some(ConflictItem::from((project_name.clone(), group.clone()))) + } + _ => None, + }) + } + + /// Return the conflict item that further restricts this scope, if any. + fn conflict(&self) -> Option> { + self.conflict.as_ref().map(ConflictItem::as_ref) + } + + fn matches(&self, env: &ResolverEnvironment) -> bool { + env.included_by_marker(self.marker) + && self + .conflict() + .is_none_or(|conflict| env.included_by_group(conflict)) + } } impl Default for ForkMap { @@ -29,15 +87,13 @@ impl Default for ForkMap { impl ForkMap { /// Associate a value with the [`Requirement`] in a given fork. pub(crate) fn add(&mut self, requirement: &Requirement, value: T) { - let entry = Entry { - value, - marker: requirement.marker, - }; - self.0 .entry(requirement.name.clone()) .or_default() - .push(entry); + .push(Entry { + value, + scope: ForkScope::from_requirement(requirement), + }); } /// Returns `true` if the map contains any values for a package that are compatible with the @@ -62,8 +118,55 @@ impl ForkMap { }; values .iter() - .filter(|entry| env.included_by_marker(entry.marker)) + .filter(|entry| entry.scope.matches(env)) .map(|entry| &entry.value) .collect() } } + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + use std::str::FromStr; + + use uv_distribution_types::RequirementSource; + use uv_normalize::{GroupName, PackageName}; + use uv_pep508::VerbatimUrl; + + use super::*; + + #[test] + fn add_scopes_non_registry_requirements_to_group_origin() { + let project_name = PackageName::from_str("workspace-root").unwrap(); + let group = GroupName::from_str("dev").unwrap(); + let package_name = PackageName::from_str("demo").unwrap(); + let conflict = ConflictItem::from((project_name.clone(), group.clone())); + let requirement = Requirement { + name: package_name.clone(), + extras: Box::default(), + groups: Box::default(), + marker: MarkerTree::TRUE, + source: RequirementSource::Directory { + install_path: PathBuf::from("/tmp/demo").into_boxed_path(), + editable: None, + r#virtual: None, + url: VerbatimUrl::parse_url("file:///tmp/demo").unwrap(), + }, + origin: Some(RequirementOrigin::Group( + PathBuf::from("pyproject.toml"), + Some(project_name), + group, + )), + }; + + let mut map = ForkMap::default(); + map.add(&requirement, ()); + + assert!(map.contains(&package_name, &ResolverEnvironment::universal(vec![]))); + + let env = ResolverEnvironment::universal(vec![]) + .filter_by_group([Err(conflict)]) + .unwrap(); + assert!(!map.contains(&package_name, &env)); + } +} diff --git a/crates/uv-resolver/src/resolver/indexes.rs b/crates/uv-resolver/src/resolver/indexes.rs index c4f955402..a4b7c7b6a 100644 --- a/crates/uv-resolver/src/resolver/indexes.rs +++ b/crates/uv-resolver/src/resolver/indexes.rs @@ -1,9 +1,7 @@ -use uv_distribution_types::{IndexMetadata, RequirementSource}; -use uv_normalize::PackageName; -use uv_pypi_types::ConflictItem; - use crate::resolver::ForkMap; use crate::{DependencyMode, Manifest, ResolverEnvironment}; +use uv_distribution_types::{IndexMetadata, RequirementSource}; +use uv_normalize::PackageName; /// A map of package names to their explicit index. /// @@ -19,13 +17,7 @@ use crate::{DependencyMode, Manifest, ResolverEnvironment}; /// /// [`Indexes`] would contain a single entry mapping `torch` to `https://download.pytorch.org/whl/cu121`. #[derive(Debug, Default, Clone)] -pub(crate) struct Indexes(ForkMap); - -#[derive(Debug, Clone)] -struct Entry { - index: IndexMetadata, - conflict: Option, -} +pub(crate) struct Indexes(ForkMap); impl Indexes { /// Determine the set of explicit, pinned indexes in the [`Manifest`]. @@ -38,16 +30,12 @@ impl Indexes { for requirement in manifest.requirements(env, dependencies) { let RequirementSource::Registry { - index: Some(index), - conflict, - .. + index: Some(index), .. } = &requirement.source else { continue; }; - let index = index.clone(); - let conflict = conflict.clone(); - indexes.add(&requirement, Entry { index, conflict }); + indexes.add(requirement.as_ref(), index.clone()); } Self(indexes) @@ -60,16 +48,6 @@ impl Indexes { /// Return the explicit index used for a package in the given fork. pub(crate) fn get(&self, name: &PackageName, env: &ResolverEnvironment) -> Vec<&IndexMetadata> { - let entries = self.0.get(name, env); - entries - .iter() - .filter(|entry| { - entry - .conflict - .as_ref() - .is_none_or(|conflict| env.included_by_group(conflict.as_ref())) - }) - .map(|entry| &entry.index) - .collect() + self.0.get(name, env) } } diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index 5509c49f0..b0d26cafa 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -50,7 +50,8 @@ use crate::manifest::Manifest; use crate::pins::FilePins; use crate::preferences::{PreferenceSource, Preferences}; use crate::pubgrub::{ - PubGrubDependency, PubGrubPackage, PubGrubPackageInner, PubGrubPriorities, PubGrubPython, + DependencySource, PubGrubDependency, PubGrubPackage, PubGrubPackageInner, PubGrubPriorities, + PubGrubPython, }; use crate::python_requirement::PythonRequirement; use crate::resolution::ResolverOutput; @@ -952,7 +953,7 @@ impl ResolverState ResolverState ResolverState ResolverState for PubGrubDependency { package: PubGrubPackage::from(PubGrubPackageInner::System(value.name)), version: Ranges::singleton(value.version), parent: None, - url: None, + source: DependencySource::Unspecified, } } } diff --git a/crates/uv-resolver/src/resolver/urls.rs b/crates/uv-resolver/src/resolver/urls.rs index eca87ef05..e1f4432fa 100644 --- a/crates/uv-resolver/src/resolver/urls.rs +++ b/crates/uv-resolver/src/resolver/urls.rs @@ -6,9 +6,10 @@ use tracing::debug; use uv_cache_key::CanonicalUrl; use uv_git::GitResolver; use uv_normalize::PackageName; -use uv_pep508::{MarkerTree, VerbatimUrl}; +use uv_pep508::VerbatimUrl; use uv_pypi_types::{ParsedDirectoryUrl, ParsedUrl, VerbatimParsedUrl}; +use crate::resolver::ForkMap; use crate::{DependencyMode, Manifest, ResolveError, ResolverEnvironment}; /// The URLs that are allowed for packages. @@ -25,7 +26,7 @@ pub(crate) struct Urls { /// URL requirements in overrides. An override URL replaces all requirements and constraints /// URLs. There can be multiple URLs for the same package as long as they are in different /// forks. - overrides: FxHashMap>, + overrides: ForkMap, /// URLs from regular requirements or from constraints. There can be multiple URLs for the same /// package as long as they are in different forks. regular: FxHashMap>, @@ -39,8 +40,7 @@ impl Urls { dependencies: DependencyMode, ) -> Self { let mut regular: FxHashMap> = FxHashMap::default(); - let mut overrides: FxHashMap> = - FxHashMap::default(); + let mut overrides = ForkMap::default(); // Add all direct regular requirements and constraints URL. for requirement in manifest.requirements_no_overrides(env, dependencies) { @@ -85,10 +85,7 @@ impl Urls { // a requirements.txt entry `./anyio`, we still use the URL. See // `allow_recursive_url_local_path_override_constraint`. regular.remove(&requirement.name); - overrides - .entry(requirement.name.clone()) - .or_default() - .push((requirement.marker, url)); + overrides.add(requirement.as_ref(), url); } Self { overrides, regular } @@ -108,16 +105,10 @@ impl Urls { url: Option<&'a VerbatimParsedUrl>, git: &'a GitResolver, ) -> Result, ResolveError> { - if let Some(override_urls) = self.get_overrides(name) { - Ok(Either::Left(Either::Left(override_urls.iter().filter_map( - |(marker, url)| { - if env.included_by_marker(*marker) { - Some(url) - } else { - None - } - }, - )))) + if self.overrides.contains_key(name) { + Ok(Either::Left(Either::Left( + self.overrides.get(name, env).into_iter(), + ))) } else if let Some(url) = url { let url = self.canonicalize_allowed_url(env, name, git, &url.verbatim, &url.parsed_url)?; @@ -129,12 +120,7 @@ impl Urls { /// Return `true` if the package has any URL (from overrides or regular requirements). pub(crate) fn any_url(&self, name: &PackageName) -> bool { - self.get_overrides(name).is_some() || self.get_regular(name).is_some() - } - - /// Return the [`VerbatimUrl`] override for the given package, if any. - fn get_overrides(&self, package: &PackageName) -> Option<&[(MarkerTree, VerbatimParsedUrl)]> { - self.overrides.get(package).map(Vec::as_slice) + self.overrides.contains_key(name) || self.get_regular(name).is_some() } /// Return the allowed [`VerbatimUrl`]s for given package from regular requirements and