diff --git a/crates/uv-resolver/src/resolver/locals.rs b/crates/uv-resolver/src/resolver/locals.rs index a515d3e9d..5ae47fcd9 100644 --- a/crates/uv-resolver/src/resolver/locals.rs +++ b/crates/uv-resolver/src/resolver/locals.rs @@ -1,5 +1,10 @@ +use std::str::FromStr; + +use either::Either; use rustc_hash::FxHashMap; +use distribution_filename::{SourceDistFilename, WheelFilename}; +use distribution_types::RemoteSource; use pep440_rs::{Operator, Version, VersionSpecifier, VersionSpecifierBuildError}; use pep508_rs::{MarkerEnvironment, VersionOrUrl}; use uv_normalize::PackageName; @@ -42,13 +47,9 @@ impl Locals { .filter(|requirement| requirement.evaluate_markers(markers, &[])), ) { - if let Some(VersionOrUrl::VersionSpecifier(specifiers)) = - requirement.version_or_url.as_ref() - { - for specifier in specifiers.iter() { - if let Some(version) = to_local(specifier) { - required.insert(requirement.name.clone(), version.clone()); - } + if let Some(version_or_url) = requirement.version_or_url.as_ref() { + for local in iter_locals(version_or_url) { + required.insert(requirement.name.clone(), local); } } } @@ -158,18 +159,48 @@ fn is_compatible(expected: &Version, provided: &Version) -> bool { } } -/// If a [`VersionSpecifier`] represents exact equality against a local version, return the local -/// version. -fn to_local(specifier: &VersionSpecifier) -> Option<&Version> { - if !matches!(specifier.operator(), Operator::Equal | Operator::ExactEqual) { - return None; - }; - - if specifier.version().local().is_empty() { - return None; +/// If a [`VersionSpecifier`] contains exact equality specifiers for a local version, returns an +/// iterator over the local versions. +fn iter_locals(version_or_url: &VersionOrUrl) -> impl Iterator + '_ { + match version_or_url { + // Extract all local versions from specifiers that require an exact version (e.g., + // `==1.0.0+local`). + VersionOrUrl::VersionSpecifier(specifiers) => Either::Left( + specifiers + .iter() + .filter(|specifier| { + matches!(specifier.operator(), Operator::Equal | Operator::ExactEqual) + }) + .filter(|specifier| !specifier.version().local().is_empty()) + .map(|specifier| specifier.version().clone()), + ), + // Exact a local version from a URL, if it includes a fully-qualified filename (e.g., + // `torch-2.2.1%2Bcu118-cp311-cp311-linux_x86_64.whl`). + VersionOrUrl::Url(url) => Either::Right( + url.filename() + .ok() + .and_then(|filename| { + if let Ok(filename) = WheelFilename::from_str(&filename) { + if filename.version.is_local() { + Some(filename.version) + } else { + None + } + } else if let Ok(filename) = + SourceDistFilename::parsed_normalized_filename(&filename) + { + if filename.version.is_local() { + Some(filename.version) + } else { + None + } + } else { + None + } + }) + .into_iter(), + ), } - - Some(specifier.version()) } #[cfg(test)] @@ -177,10 +208,53 @@ mod tests { use std::str::FromStr; use anyhow::Result; + use url::Url; - use pep440_rs::{Operator, Version, VersionSpecifier}; + use pep440_rs::{Operator, Version, VersionSpecifier, VersionSpecifiers}; + use pep508_rs::{VerbatimUrl, VersionOrUrl}; - use super::Locals; + use crate::resolver::locals::{iter_locals, Locals}; + + #[test] + fn extract_locals() -> Result<()> { + // Extract from a source distribution in a URL. + let version_or_url = VersionOrUrl::Url(VerbatimUrl::from_url(Url::parse( + "https://example.com/foo-1.0.0+local.tar.gz", + )?)); + let locals: Vec<_> = iter_locals(&version_or_url).collect(); + assert_eq!(locals, vec![Version::from_str("1.0.0+local")?]); + + // Extract from a wheel in a URL. + let version_or_url = VersionOrUrl::Url(VerbatimUrl::from_url(Url::parse( + "https://example.com/foo-1.0.0+local-cp39-cp39-linux_x86_64.whl", + )?)); + let locals: Vec<_> = iter_locals(&version_or_url).collect(); + assert_eq!(locals, vec![Version::from_str("1.0.0+local")?]); + + // Don't extract anything if the URL is opaque. + let version_or_url = VersionOrUrl::Url(VerbatimUrl::from_url(Url::parse( + "git+https://example.com/foo/bar", + )?)); + let locals: Vec<_> = iter_locals(&version_or_url).collect(); + assert!(locals.is_empty()); + + // Extract from `==` specifiers. + let version_or_url = VersionOrUrl::VersionSpecifier(VersionSpecifiers::from_iter([ + VersionSpecifier::from_version(Operator::GreaterThan, Version::from_str("1.0.0")?)?, + VersionSpecifier::from_version(Operator::Equal, Version::from_str("1.0.0+local")?)?, + ])); + let locals: Vec<_> = iter_locals(&version_or_url).collect(); + assert_eq!(locals, vec![Version::from_str("1.0.0+local")?]); + + // Ignore other specifiers. + let version_or_url = VersionOrUrl::VersionSpecifier(VersionSpecifiers::from_iter([ + VersionSpecifier::from_version(Operator::NotEqual, Version::from_str("1.0.0+local")?)?, + ])); + let locals: Vec<_> = iter_locals(&version_or_url).collect(); + assert!(locals.is_empty()); + + Ok(()) + } #[test] fn map_version() -> Result<()> {