Skip to content

Commit

Permalink
Extract supported architectures from wheel tags (#10179)
Browse files Browse the repository at this point in the history 
## Summary

This PR extends #10046 to also handle architectures, which allows us to
correctly include `2.5.1` on the `cu124` index for ARM Linux.

Closes #9655.
  • Loading branch information
@charliermarsh
charliermarsh authored Jan 4, 2025
1 parent fbe6f1e commit 833519d
Show file tree
Hide file tree
Showing 4 changed files with 409 additions and 108 deletions.
278 changes: 266 additions & 12 deletions crates/uv-distribution-types/src/prioritized_distribution.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use std::fmt::{Display, Formatter};
use uv_distribution_filename::{BuildTag, WheelFilename};

use tracing::debug;

use uv_distribution_filename::{BuildTag, WheelFilename};
use uv_pep440::VersionSpecifiers;
use uv_pep508::{MarkerExpression, MarkerOperator, MarkerTree, MarkerValueString};
use uv_platform_tags::{IncompatibleTag, TagPriority};
Expand Down Expand Up @@ -634,39 +636,291 @@ impl IncompatibleWheel {
}

/// Given a wheel filename, determine the set of supported platforms, in terms of their markers.
///
/// This is roughly the inverse of platform tag generation: given a tag, we want to infer the
/// supported platforms (rather than generating the supported tags from a given platform).
pub fn implied_markers(filename: &WheelFilename) -> MarkerTree {
let mut marker = MarkerTree::FALSE;
for platform_tag in &filename.platform_tag {
match platform_tag.as_str() {
"any" => {
return MarkerTree::TRUE;
}
tag if tag.starts_with("win") => {
marker.or(MarkerTree::expression(MarkerExpression::String {

// Windows
"win32" => {
let mut tag_marker = MarkerTree::expression(MarkerExpression::String {
key: MarkerValueString::SysPlatform,
operator: MarkerOperator::Equal,
value: "win32".to_string(),
});
tag_marker.and(MarkerTree::expression(MarkerExpression::String {
key: MarkerValueString::PlatformMachine,
operator: MarkerOperator::Equal,
value: "x86".to_string(),
}));
marker.or(tag_marker);
}
tag if tag.starts_with("macosx") => {
marker.or(MarkerTree::expression(MarkerExpression::String {
"win_amd64" => {
let mut tag_marker = MarkerTree::expression(MarkerExpression::String {
key: MarkerValueString::SysPlatform,
operator: MarkerOperator::Equal,
value: "darwin".to_string(),
value: "win32".to_string(),
});
tag_marker.and(MarkerTree::expression(MarkerExpression::String {
key: MarkerValueString::PlatformMachine,
operator: MarkerOperator::Equal,
value: "x86_64".to_string(),
}));
marker.or(tag_marker);
}
"win_arm64" => {
let mut tag_marker = MarkerTree::expression(MarkerExpression::String {
key: MarkerValueString::SysPlatform,
operator: MarkerOperator::Equal,
value: "win32".to_string(),
});
tag_marker.and(MarkerTree::expression(MarkerExpression::String {
key: MarkerValueString::PlatformMachine,
operator: MarkerOperator::Equal,
value: "arm64".to_string(),
}));
marker.or(tag_marker);
}

// macOS
tag if tag.starts_with("macosx_") => {
let mut tag_marker = MarkerTree::expression(MarkerExpression::String {
key: MarkerValueString::SysPlatform,
operator: MarkerOperator::Equal,
value: "darwin".to_string(),
});

// Parse the macOS version from the tag.
//
// For example, given `macosx_10_9_x86_64`, infer `10.9`, followed by `x86_64`.
//
// If at any point we fail to parse, we assume the tag is invalid and skip it.
let mut parts = tag.splitn(4, '_');

// Skip the "macosx_" prefix.
if parts.next().is_none_or(|part| part != "macosx") {
debug!("Failed to parse macOS prefix from tag: {tag}");
continue;
}

// Skip the major and minor version numbers.
if parts
.next()
.and_then(|part| part.parse::<u16>().ok())
.is_none()
{
debug!("Failed to parse macOS major version from tag: {tag}");
continue;
};
if parts
.next()
.and_then(|part| part.parse::<u16>().ok())
.is_none()
{
debug!("Failed to parse macOS minor version from tag: {tag}");
continue;
};

// Extract the architecture from the end of the tag.
let Some(arch) = parts.next() else {
debug!("Failed to parse macOS architecture from tag: {tag}");
continue;
};

// Extract the architecture from the end of the tag.
let mut arch_marker = MarkerTree::FALSE;
let supported_architectures = match arch {
"universal" => {
// Allow any of: "x86_64", "i386", "ppc64", "ppc", "intel"
["x86_64", "i386", "ppc64", "ppc", "intel"].iter()
}
"universal2" => {
// Allow any of: "x86_64", "arm64"
["x86_64", "arm64"].iter()
}
"intel" => {
// Allow any of: "x86_64", "i386"
["x86_64", "i386"].iter()
}
"x86_64" => {
// Allow only "x86_64"
["x86_64"].iter()
}
"arm64" => {
// Allow only "arm64"
["arm64"].iter()
}
"ppc64" => {
// Allow only "ppc64"
["ppc64"].iter()
}
"ppc" => {
// Allow only "ppc"
["ppc"].iter()
}
"i386" => {
// Allow only "i386"
["i386"].iter()
}
_ => {
debug!("Unknown macOS architecture in wheel tag: {tag}");
continue;
}
};
for arch in supported_architectures {
arch_marker.or(MarkerTree::expression(MarkerExpression::String {
key: MarkerValueString::PlatformMachine,
operator: MarkerOperator::Equal,
value: (*arch).to_string(),
}));
}
tag_marker.and(arch_marker);

marker.or(tag_marker);
}
tag if tag.starts_with("manylinux")
|| tag.starts_with("musllinux")
|| tag.starts_with("linux") =>
{
marker.or(MarkerTree::expression(MarkerExpression::String {

// Linux
tag => {
let mut tag_marker = MarkerTree::expression(MarkerExpression::String {
key: MarkerValueString::SysPlatform,
operator: MarkerOperator::Equal,
value: "linux".to_string(),
});

// Parse the architecture from the tag.
let arch = if let Some(arch) = tag.strip_prefix("linux_") {
arch
} else if let Some(arch) = tag.strip_prefix("manylinux1_") {
arch
} else if let Some(arch) = tag.strip_prefix("manylinux2010_") {
arch
} else if let Some(arch) = tag.strip_prefix("manylinux2014_") {
arch
} else if let Some(arch) = tag.strip_prefix("musllinux_") {
// Skip over the version tags (e.g., given `musllinux_1_2`, skip over `1` and `2`).
let mut parts = arch.splitn(3, '_');
if parts
.next()
.and_then(|part| part.parse::<u16>().ok())
.is_none()
{
debug!("Failed to parse musllinux major version from tag: {tag}");
continue;
};
if parts
.next()
.and_then(|part| part.parse::<u16>().ok())
.is_none()
{
debug!("Failed to parse musllinux minor version from tag: {tag}");
continue;
};
let Some(arch) = parts.next() else {
debug!("Failed to parse musllinux architecture from tag: {tag}");
continue;
};
arch
} else if let Some(arch) = tag.strip_prefix("manylinux_") {
// Skip over the version tags (e.g., given `manylinux_2_17`, skip over `2` and `17`).
let mut parts = arch.splitn(3, '_');
if parts
.next()
.and_then(|part| part.parse::<u16>().ok())
.is_none()
{
debug!("Failed to parse manylinux major version from tag: {tag}");
continue;
};
if parts
.next()
.and_then(|part| part.parse::<u16>().ok())
.is_none()
{
debug!("Failed to parse manylinux minor version from tag: {tag}");
continue;
};
let Some(arch) = parts.next() else {
debug!("Failed to parse manylinux architecture from tag: {tag}");
continue;
};
arch
} else {
continue;
};
tag_marker.and(MarkerTree::expression(MarkerExpression::String {
key: MarkerValueString::PlatformMachine,
operator: MarkerOperator::Equal,
value: arch.to_string(),
}));

marker.or(tag_marker);
}
_ => {}
}
}
marker
}

#[cfg(test)]
mod tests {
use std::str::FromStr;

use super::*;

#[track_caller]
fn assert_markers(filename: &str, expected: &str) {
let filename = WheelFilename::from_str(filename).unwrap();
assert_eq!(
implied_markers(&filename),
expected.parse::<MarkerTree>().unwrap()
);
}

#[test]
fn test_implied_markers() {
let filename = WheelFilename::from_str("example-1.0-py3-none-any.whl").unwrap();
assert_eq!(implied_markers(&filename), MarkerTree::TRUE);

assert_markers(
"example-1.0-cp310-cp310-win32.whl",
"sys_platform == 'win32' and platform_machine == 'x86'",
);
assert_markers(
"numpy-2.2.1-cp313-cp313t-win_amd64.whl",
"sys_platform == 'win32' and platform_machine == 'x86_64'",
);
assert_markers(
"numpy-2.2.1-cp313-cp313t-win_arm64.whl",
"sys_platform == 'win32' and platform_machine == 'arm64'",
);
assert_markers(
"numpy-2.2.1-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl",
"sys_platform == 'linux' and platform_machine == 'aarch64'",
);
assert_markers(
"numpy-2.2.1-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl",
"sys_platform == 'linux' and platform_machine == 'x86_64'",
);
assert_markers(
"numpy-2.2.1-cp312-cp312-musllinux_1_2_aarch64.whl",
"sys_platform == 'linux' and platform_machine == 'aarch64'",
);
assert_markers(
"numpy-2.2.1-cp310-cp310-macosx_14_0_x86_64.whl",
"sys_platform == 'darwin' and platform_machine == 'x86_64'",
);
assert_markers(
"numpy-2.2.1-cp310-cp310-macosx_10_9_x86_64.whl",
"sys_platform == 'darwin' and platform_machine == 'x86_64'",
);
assert_markers(
"numpy-2.2.1-cp310-cp310-macosx_11_0_arm64.whl",
"sys_platform == 'darwin' and platform_machine == 'arm64'",
);
}
}
25 changes: 23 additions & 2 deletions crates/uv-resolver/src/resolver/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ use uv_distribution_types::{
use uv_git::GitResolver;
use uv_normalize::{ExtraName, GroupName, PackageName};
use uv_pep440::{release_specifiers_to_ranges, Version, VersionSpecifiers, MIN_VERSION};
use uv_pep508::MarkerTree;
use uv_pep508::{MarkerExpression, MarkerOperator, MarkerTree, MarkerValueString};
use uv_platform_tags::Tags;
use uv_pypi_types::{ConflictItem, ConflictItemRef, Conflicts, Requirement, VerbatimParsedUrl};
use uv_types::{BuildContext, HashStrategy, InstalledPackagesProvider};
Expand Down Expand Up @@ -1372,7 +1372,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
};

// ...and the non-local version has greater platform support...
let remainder = {
let mut remainder = {
let mut remainder = base_dist.implied_markers();
remainder.and(dist.implied_markers().negate());
remainder
Expand Down Expand Up @@ -1415,6 +1415,27 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
)));
}

// If the implied markers includes _some_ macOS environments, but the remainder doesn't,
// then we can extend the implied markers to include _all_ macOS environments. Same goes for
// Linux and Windows.
//
// The idea here is that the base version could support (e.g.) ARM macOS, but not Intel
// macOS. But if _neither_ version supports Intel macOS, we'd rather use `sys_platform == 'darwin'`
// instead of `sys_platform == 'darwin' and platform_machine == 'arm64'`, since it's much
// simpler, and _neither_ version will succeed with Intel macOS anyway.
for sys_platform in &["darwin", "linux", "win32"] {
let sys_platform = MarkerTree::expression(MarkerExpression::String {
key: MarkerValueString::SysPlatform,
operator: MarkerOperator::Equal,
value: (*sys_platform).to_string(),
});
if dist.implied_markers().is_disjoint(sys_platform)
&& !remainder.is_disjoint(sys_platform)
{
remainder.or(sys_platform);
}
}

// Otherwise, we need to fork.
let Some((base_env, local_env)) = fork_version_by_marker(env, remainder) else {
return Ok(None);
Expand Down
Loading

0 comments on commit 833519d

Please sign in to comment.