Reduce `WheelFilename` to 48 bytes (#10583)

## Summary

Based on some advice from @konstin.
This commit is contained in:
Charlie Marsh 2025-01-14 09:49:17 -05:00 committed by GitHub
parent e3f6b9c5f3
commit 24a5920739
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
16 changed files with 379 additions and 227 deletions

View File

@ -10,7 +10,7 @@ use tracing::{debug, trace};
use walkdir::WalkDir;
use zip::{CompressionMethod, ZipWriter};
use uv_distribution_filename::{TagSet, WheelFilename};
use uv_distribution_filename::WheelFilename;
use uv_fs::Simplified;
use uv_globfilter::{parse_portable_glob, GlobDirFilter};
use uv_platform_tags::{AbiTag, LanguageTag, PlatformTag};
@ -33,17 +33,16 @@ pub fn build_wheel(
}
crate::check_metadata_directory(source_tree, metadata_directory, &pyproject_toml)?;
let filename = WheelFilename {
name: pyproject_toml.name().clone(),
version: pyproject_toml.version().clone(),
build_tag: None,
python_tag: TagSet::from_slice(&[LanguageTag::Python {
let filename = WheelFilename::new(
pyproject_toml.name().clone(),
pyproject_toml.version().clone(),
LanguageTag::Python {
major: 3,
minor: None,
}]),
abi_tag: TagSet::from_buf([AbiTag::None]),
platform_tag: TagSet::from_buf([PlatformTag::Any]),
};
},
AbiTag::None,
PlatformTag::Any,
);
let wheel_path = wheel_dir.join(filename.to_string());
debug!("Writing wheel at {}", wheel_path.user_display());
@ -71,17 +70,16 @@ pub fn list_wheel(
warn_user_once!("{warning}");
}
let filename = WheelFilename {
name: pyproject_toml.name().clone(),
version: pyproject_toml.version().clone(),
build_tag: None,
python_tag: TagSet::from_slice(&[LanguageTag::Python {
let filename = WheelFilename::new(
pyproject_toml.name().clone(),
pyproject_toml.version().clone(),
LanguageTag::Python {
major: 3,
minor: None,
}]),
abi_tag: TagSet::from_buf([AbiTag::None]),
platform_tag: TagSet::from_buf([PlatformTag::Any]),
};
},
AbiTag::None,
PlatformTag::Any,
);
let mut files = FileList::new();
let writer = ListWriter::new(&mut files);
@ -253,17 +251,16 @@ pub fn build_editable(
crate::check_metadata_directory(source_tree, metadata_directory, &pyproject_toml)?;
let filename = WheelFilename {
name: pyproject_toml.name().clone(),
version: pyproject_toml.version().clone(),
build_tag: None,
python_tag: TagSet::from_slice(&[LanguageTag::Python {
let filename = WheelFilename::new(
pyproject_toml.name().clone(),
pyproject_toml.version().clone(),
LanguageTag::Python {
major: 3,
minor: None,
}]),
abi_tag: TagSet::from_buf([AbiTag::None]),
platform_tag: TagSet::from_buf([PlatformTag::Any]),
};
},
AbiTag::None,
PlatformTag::Any,
);
let wheel_path = wheel_dir.join(filename.to_string());
debug!("Writing wheel at {}", wheel_path.user_display());
@ -308,17 +305,16 @@ pub fn metadata(
warn_user_once!("{warning}");
}
let filename = WheelFilename {
name: pyproject_toml.name().clone(),
version: pyproject_toml.version().clone(),
build_tag: None,
python_tag: TagSet::from_slice(&[LanguageTag::Python {
let filename = WheelFilename::new(
pyproject_toml.name().clone(),
pyproject_toml.version().clone(),
LanguageTag::Python {
major: 3,
minor: None,
}]),
abi_tag: TagSet::from_buf([AbiTag::None]),
platform_tag: TagSet::from_buf([PlatformTag::Any]),
};
},
AbiTag::None,
PlatformTag::Any,
);
debug!(
"Writing metadata files to {}",
@ -568,9 +564,9 @@ fn wheel_info(filename: &WheelFilename, uv_version: &str) -> String {
("Generator", format!("uv {uv_version}")),
("Root-Is-Purelib", "true".to_string()),
];
for python_tag in &filename.python_tag {
for abi_tag in &filename.abi_tag {
for platform_tag in &filename.platform_tag {
for python_tag in filename.python_tags() {
for abi_tag in filename.abi_tags() {
for platform_tag in filename.platform_tags() {
wheel_info.push(("Tag", format!("{python_tag}-{abi_tag}-{platform_tag}")));
}
}
@ -765,29 +761,21 @@ mod test {
#[test]
fn test_wheel() {
let filename = WheelFilename {
name: PackageName::from_str("foo").unwrap(),
version: Version::from_str("1.2.3").unwrap(),
build_tag: None,
python_tag: TagSet::from_slice(&[
LanguageTag::Python {
major: 2,
minor: None,
},
LanguageTag::Python {
major: 3,
minor: None,
},
]),
abi_tag: TagSet::from_buf([AbiTag::None]),
platform_tag: TagSet::from_buf([PlatformTag::Any]),
};
let filename = WheelFilename::new(
PackageName::from_str("foo").unwrap(),
Version::from_str("1.2.3").unwrap(),
LanguageTag::Python {
major: 3,
minor: None,
},
AbiTag::None,
PlatformTag::Any,
);
assert_snapshot!(wheel_info(&filename, "1.0.0+test"), @r"
Wheel-Version: 1.0
Generator: uv 1.0.0+test
Root-Is-Purelib: true
Tag: py2-none-any
Tag: py3-none-any
");
}

View File

@ -99,6 +99,6 @@ mod tests {
#[test]
fn wheel_filename_size() {
assert_eq!(size_of::<WheelFilename>(), 128);
assert_eq!(size_of::<WheelFilename>(), 48);
}
}

View File

@ -8,23 +8,27 @@ Ok(
"foo",
),
version: "1.2.3",
build_tag: Some(
BuildTag(
202206090410,
None,
),
),
python_tag: [
Python {
major: 3,
minor: None,
tags: Large {
large: WheelTagLarge {
build_tag: Some(
BuildTag(
202206090410,
None,
),
),
python_tag: [
Python {
major: 3,
minor: None,
},
],
abi_tag: [
None,
],
platform_tag: [
Any,
],
},
],
abi_tag: [
None,
],
platform_tag: [
Any,
],
},
},
)

View File

@ -8,33 +8,37 @@ Ok(
"foo",
),
version: "1.2.3",
build_tag: None,
python_tag: [
CPython {
python_version: (
3,
11,
),
tags: Large {
large: WheelTagLarge {
build_tag: None,
python_tag: [
CPython {
python_version: (
3,
11,
),
},
],
abi_tag: [
CPython {
gil_disabled: false,
python_version: (
3,
11,
),
},
],
platform_tag: [
Manylinux {
major: 2,
minor: 17,
arch: X86_64,
},
Manylinux2014 {
arch: X86_64,
},
],
},
],
abi_tag: [
CPython {
gil_disabled: false,
python_version: (
3,
11,
),
},
],
platform_tag: [
Manylinux {
major: 2,
minor: 17,
arch: X86_64,
},
Manylinux2014 {
arch: X86_64,
},
],
},
},
)

View File

@ -8,18 +8,15 @@ Ok(
"foo",
),
version: "1.2.3",
build_tag: None,
python_tag: [
Python {
major: 3,
minor: None,
tags: Small {
small: WheelTagSmall {
python_tag: Python {
major: 3,
minor: None,
},
abi_tag: None,
platform_tag: Any,
},
],
abi_tag: [
None,
],
platform_tag: [
Any,
],
},
},
)

View File

@ -1,7 +1,7 @@
use std::fmt::{Display, Formatter};
use std::str::FromStr;
use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
use std::str::FromStr;
use thiserror::Error;
use url::Url;
@ -14,12 +14,6 @@ use uv_platform_tags::{
use crate::{BuildTag, BuildTagError};
/// A [`SmallVec`] type for storing tags.
///
/// Wheels tend to include a single language, ABI, and platform tag, so we use a [`SmallVec`] with a
/// capacity of 1 to optimize for this common case.
pub type TagSet<T> = smallvec::SmallVec<[T; 1]>;
#[derive(
Debug,
Clone,
@ -36,10 +30,7 @@ pub type TagSet<T> = smallvec::SmallVec<[T; 1]>;
pub struct WheelFilename {
pub name: PackageName,
pub version: Version,
pub build_tag: Option<BuildTag>,
pub python_tag: TagSet<LanguageTag>,
pub abi_tag: TagSet<AbiTag>,
pub platform_tag: TagSet<PlatformTag>,
tags: WheelTag,
}
impl FromStr for WheelFilename {
@ -63,20 +54,41 @@ impl Display for WheelFilename {
"{}-{}-{}.whl",
self.name.as_dist_info_name(),
self.version,
self.get_tag()
self.tags,
)
}
}
impl WheelFilename {
/// Create a [`WheelFilename`] from its components.
pub fn new(
name: PackageName,
version: Version,
python_tag: LanguageTag,
abi_tag: AbiTag,
platform_tag: PlatformTag,
) -> Self {
Self {
name,
version,
tags: WheelTag::Small {
small: WheelTagSmall {
python_tag,
abi_tag,
platform_tag,
},
},
}
}
/// Returns `true` if the wheel is compatible with the given tags.
pub fn is_compatible(&self, compatible_tags: &Tags) -> bool {
compatible_tags.is_compatible(&self.python_tag, &self.abi_tag, &self.platform_tag)
compatible_tags.is_compatible(self.python_tags(), self.abi_tags(), self.platform_tags())
}
/// Return the [`TagCompatibility`] of the wheel with the given tags
pub fn compatibility(&self, compatible_tags: &Tags) -> TagCompatibility {
compatible_tags.compatibility(&self.python_tag, &self.abi_tag, &self.platform_tag)
compatible_tags.compatibility(self.python_tags(), self.abi_tags(), self.platform_tags())
}
/// The wheel filename without the extension.
@ -85,45 +97,47 @@ impl WheelFilename {
"{}-{}-{}",
self.name.as_dist_info_name(),
self.version,
self.get_tag()
self.tags
)
}
/// Return the wheel's Python tags.
pub fn python_tags(&self) -> &[LanguageTag] {
match &self.tags {
WheelTag::Small { small } => std::slice::from_ref(&small.python_tag),
WheelTag::Large { large } => large.python_tag.as_slice(),
}
}
/// Return the wheel's ABI tags.
pub fn abi_tags(&self) -> &[AbiTag] {
match &self.tags {
WheelTag::Small { small } => std::slice::from_ref(&small.abi_tag),
WheelTag::Large { large } => large.abi_tag.as_slice(),
}
}
/// Return the wheel's platform tags.
pub fn platform_tags(&self) -> &[PlatformTag] {
match &self.tags {
WheelTag::Small { small } => std::slice::from_ref(&small.platform_tag),
WheelTag::Large { large } => large.platform_tag.as_slice(),
}
}
/// Return the wheel's build tag, if present.
pub fn build_tag(&self) -> Option<&BuildTag> {
match &self.tags {
WheelTag::Small { .. } => None,
WheelTag::Large { large } => large.build_tag.as_ref(),
}
}
/// Parse a wheel filename from the stem (e.g., `foo-1.2.3-py3-none-any`).
pub fn from_stem(stem: &str) -> Result<Self, WheelFilenameError> {
Self::parse(stem, stem)
}
/// Get the tag for this wheel.
fn get_tag(&self) -> String {
if let ([python_tag], [abi_tag], [platform_tag]) = (
self.python_tag.as_slice(),
self.abi_tag.as_slice(),
self.platform_tag.as_slice(),
) {
format!("{python_tag}-{abi_tag}-{platform_tag}",)
} else {
format!(
"{}-{}-{}",
self.python_tag
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join("."),
self.abi_tag
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join("."),
self.platform_tag
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join("."),
)
}
}
/// Parse a wheel filename from the stem (e.g., `foo-1.2.3-py3-none-any`).
///
/// The originating `filename` is used for high-fidelity error messages.
@ -202,25 +216,58 @@ impl WheelFilename {
.map_err(|err| WheelFilenameError::InvalidBuildTag(filename.to_string(), err))
})
.transpose()?;
let tags = if build_tag.is_some()
|| python_tag.contains('.')
|| abi_tag.contains('.')
|| platform_tag.contains('.')
{
WheelTag::Large {
large: Box::new(WheelTagLarge {
build_tag,
python_tag: python_tag
.split('.')
.map(LanguageTag::from_str)
.collect::<Result<_, _>>()
.map_err(|err| {
WheelFilenameError::InvalidLanguageTag(filename.to_string(), err)
})?,
abi_tag: abi_tag
.split('.')
.map(AbiTag::from_str)
.collect::<Result<_, _>>()
.map_err(|err| {
WheelFilenameError::InvalidAbiTag(filename.to_string(), err)
})?,
platform_tag: platform_tag
.split('.')
.map(PlatformTag::from_str)
.collect::<Result<_, _>>()
.map_err(|err| {
WheelFilenameError::InvalidPlatformTag(filename.to_string(), err)
})?,
}),
}
} else {
WheelTag::Small {
small: WheelTagSmall {
python_tag: LanguageTag::from_str(python_tag).map_err(|err| {
WheelFilenameError::InvalidLanguageTag(filename.to_string(), err)
})?,
abi_tag: AbiTag::from_str(abi_tag).map_err(|err| {
WheelFilenameError::InvalidAbiTag(filename.to_string(), err)
})?,
platform_tag: PlatformTag::from_str(platform_tag).map_err(|err| {
WheelFilenameError::InvalidPlatformTag(filename.to_string(), err)
})?,
},
}
};
Ok(Self {
name,
version,
build_tag,
python_tag: python_tag
.split('.')
.map(LanguageTag::from_str)
.collect::<Result<_, _>>()
.map_err(|err| WheelFilenameError::InvalidLanguageTag(filename.to_string(), err))?,
abi_tag: abi_tag
.split('.')
.map(AbiTag::from_str)
.collect::<Result<_, _>>()
.map_err(|err| WheelFilenameError::InvalidAbiTag(filename.to_string(), err))?,
platform_tag: platform_tag
.split('.')
.map(PlatformTag::from_str)
.collect::<Result<_, _>>()
.map_err(|err| WheelFilenameError::InvalidPlatformTag(filename.to_string(), err))?,
tags,
})
}
}
@ -267,6 +314,121 @@ impl Serialize for WheelFilename {
}
}
/// A [`SmallVec`] type for storing tags.
///
/// Wheels tend to include a single language, ABI, and platform tag, so we use a [`SmallVec`] with a
/// capacity of 1 to optimize for this common case.
pub type TagSet<T> = smallvec::SmallVec<[T; 3]>;
/// The portion of the wheel filename following the name and version: the optional build tag, along
/// with the Python tag(s), ABI tag(s), and platform tag(s).
///
/// Most wheels consist of a single Python, ABI, and platform tag (and no build tag). We represent
/// such wheels with [`WheelTagSmall`], a variant with a smaller memory footprint and (generally)
/// zero allocations. The [`WheelTagLarge`] variant is used for wheels with multiple tags and/or a
/// build tag.
#[derive(
Debug,
Clone,
Eq,
PartialEq,
Ord,
PartialOrd,
Hash,
rkyv::Archive,
rkyv::Deserialize,
rkyv::Serialize,
)]
#[rkyv(derive(Debug))]
enum WheelTag {
Small { small: WheelTagSmall },
Large { large: Box<WheelTagLarge> },
}
impl Display for WheelTag {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Small { small } => write!(f, "{small}"),
Self::Large { large } => write!(f, "{large}"),
}
}
}
#[derive(
Debug,
Clone,
Eq,
PartialEq,
Ord,
PartialOrd,
Hash,
rkyv::Archive,
rkyv::Deserialize,
rkyv::Serialize,
)]
#[rkyv(derive(Debug))]
#[allow(clippy::struct_field_names)]
struct WheelTagSmall {
python_tag: LanguageTag,
abi_tag: AbiTag,
platform_tag: PlatformTag,
}
impl Display for WheelTagSmall {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}-{}-{}",
self.python_tag, self.abi_tag, self.platform_tag
)
}
}
#[derive(
Debug,
Clone,
Eq,
PartialEq,
Ord,
PartialOrd,
Hash,
rkyv::Archive,
rkyv::Deserialize,
rkyv::Serialize,
)]
#[rkyv(derive(Debug))]
#[allow(clippy::struct_field_names)]
pub struct WheelTagLarge {
build_tag: Option<BuildTag>,
python_tag: TagSet<LanguageTag>,
abi_tag: TagSet<AbiTag>,
platform_tag: TagSet<PlatformTag>,
}
impl Display for WheelTagLarge {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}-{}-{}",
self.python_tag
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join("."),
self.abi_tag
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join("."),
self.platform_tag
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join("."),
)
}
}
#[derive(Error, Debug)]
pub enum WheelFilenameError {
#[error("The wheel filename \"{0}\" is invalid: {1}")]
@ -283,6 +445,12 @@ pub enum WheelFilenameError {
InvalidAbiTag(String, ParseAbiTagError),
#[error("The wheel filename \"{0}\" has an invalid platform tag: {1}")]
InvalidPlatformTag(String, ParsePlatformTagError),
#[error("The wheel filename \"{0}\" is missing a language tag")]
MissingLanguageTag(String),
#[error("The wheel filename \"{0}\" is missing an ABI tag")]
MissingAbiTag(String),
#[error("The wheel filename \"{0}\" is missing a platform tag")]
MissingPlatformTag(String),
}
#[cfg(test)]

View File

@ -1343,8 +1343,8 @@ mod test {
/// Ensure that we don't accidentally grow the `Dist` sizes.
#[test]
fn dist_size() {
assert!(size_of::<Dist>() <= 272, "{}", size_of::<Dist>());
assert!(size_of::<BuiltDist>() <= 272, "{}", size_of::<BuiltDist>());
assert!(size_of::<Dist>() <= 248, "{}", size_of::<Dist>());
assert!(size_of::<BuiltDist>() <= 216, "{}", size_of::<BuiltDist>());
assert!(
size_of::<SourceDist>() <= 248,
"{}",

View File

@ -527,7 +527,7 @@ impl PrioritizedDist {
self.0
.wheels
.iter()
.flat_map(|(wheel, _)| wheel.filename.python_tag.iter().copied())
.flat_map(|(wheel, _)| wheel.filename.python_tags().iter().copied())
.collect()
}
@ -536,7 +536,7 @@ impl PrioritizedDist {
self.0
.wheels
.iter()
.flat_map(|(wheel, _)| wheel.filename.abi_tag.iter().copied())
.flat_map(|(wheel, _)| wheel.filename.abi_tags().iter().copied())
.collect()
}
@ -545,10 +545,10 @@ impl PrioritizedDist {
pub fn platform_tags<'a>(&'a self, tags: &'a Tags) -> BTreeSet<&'a PlatformTag> {
let mut candidates = BTreeSet::new();
for (wheel, _) in &self.0.wheels {
for wheel_py in &wheel.filename.python_tag {
for wheel_abi in &wheel.filename.abi_tag {
for wheel_py in wheel.filename.python_tags() {
for wheel_abi in wheel.filename.abi_tags() {
if tags.is_compatible_abi(*wheel_py, *wheel_abi) {
candidates.extend(wheel.filename.platform_tag.iter());
candidates.extend(wheel.filename.platform_tags().iter());
}
}
}
@ -724,7 +724,7 @@ impl IncompatibleWheel {
/// 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 {
for platform_tag in filename.platform_tags() {
match platform_tag {
PlatformTag::Any => {
return MarkerTree::TRUE;

View File

@ -674,7 +674,7 @@ async fn form_metadata(
];
if let DistFilename::WheelFilename(wheel) = filename {
form_metadata.push(("pyversion", wheel.python_tag.iter().join(".")));
form_metadata.push(("pyversion", wheel.python_tags().iter().join(".")));
} else {
form_metadata.push(("pyversion", "source".to_string()));
}

View File

@ -181,7 +181,7 @@ impl FlatIndex {
};
// Break ties with the build tag.
let build_tag = filename.build_tag.clone();
let build_tag = filename.build_tag().cloned();
WheelCompatibility::Compatible(hash, priority, build_tag)
}

View File

@ -291,7 +291,7 @@ impl Lock {
// `(A ∩ (B ∩ C) = ∅) => ((A ∩ B = ∅) or (A ∩ C = ∅))`
// a single disjointness check with the intersection is sufficient, so we have one
// constant per platform.
let platform_tags = &wheel.filename.platform_tag;
let platform_tags = wheel.filename.platform_tags();
if platform_tags
.iter()
.all(uv_platform_tags::PlatformTag::is_linux)
@ -2275,7 +2275,7 @@ impl Package {
else {
continue;
};
let build_tag = wheel.filename.build_tag.as_ref();
let build_tag = wheel.filename.build_tag();
let wheel_priority = (tag_priority, build_tag);
match best {
None => {

View File

@ -66,19 +66,16 @@ Ok(
"anyio",
),
version: "4.3.0",
build_tag: None,
python_tag: [
Python {
major: 3,
minor: None,
tags: Small {
small: WheelTagSmall {
python_tag: Python {
major: 3,
minor: None,
},
abi_tag: None,
platform_tag: Any,
},
],
abi_tag: [
None,
],
platform_tag: [
Any,
],
},
},
},
],

View File

@ -73,19 +73,16 @@ Ok(
"anyio",
),
version: "4.3.0",
build_tag: None,
python_tag: [
Python {
major: 3,
minor: None,
tags: Small {
small: WheelTagSmall {
python_tag: Python {
major: 3,
minor: None,
},
abi_tag: None,
platform_tag: Any,
},
],
abi_tag: [
None,
],
platform_tag: [
Any,
],
},
},
},
],

View File

@ -69,19 +69,16 @@ Ok(
"anyio",
),
version: "4.3.0",
build_tag: None,
python_tag: [
Python {
major: 3,
minor: None,
tags: Small {
small: WheelTagSmall {
python_tag: Python {
major: 3,
minor: None,
},
abi_tag: None,
platform_tag: Any,
},
],
abi_tag: [
None,
],
platform_tag: [
Any,
],
},
},
},
],

View File

@ -380,12 +380,12 @@ impl RequiresPython {
/// It is meant to filter out clearly unusable wheels with perfect specificity and acceptable
/// sensitivity, we return `true` if the tags are unknown.
pub fn matches_wheel_tag(&self, wheel: &WheelFilename) -> bool {
wheel.abi_tag.iter().any(|abi_tag| {
wheel.abi_tags().iter().any(|abi_tag| {
if *abi_tag == AbiTag::Abi3 {
// Universal tags are allowed.
true
} else if *abi_tag == AbiTag::None {
wheel.python_tag.iter().any(|python_tag| {
wheel.python_tags().iter().any(|python_tag| {
// Remove `py2-none-any` and `py27-none-any` and analogous `cp` and `pp` tags.
if matches!(
python_tag,

View File

@ -555,7 +555,7 @@ impl VersionMapLazy {
};
// Break ties with the build tag.
let build_tag = filename.build_tag.clone();
let build_tag = filename.build_tag().cloned();
WheelCompatibility::Compatible(hash, priority, build_tag)
}