Skip to content

Commit

Permalink
Store unsupported tags in wheel filename
Browse files Browse the repository at this point in the history
  • Loading branch information
charliermarsh committed Jan 16, 2025
1 parent ee6ba41 commit 7740d8d
Show file tree
Hide file tree
Showing 12 changed files with 271 additions and 305 deletions.
9 changes: 3 additions & 6 deletions crates/uv-distribution-filename/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,15 @@ pub use build_tag::{BuildTag, BuildTagError};
pub use egg::{EggInfoFilename, EggInfoFilenameError};
pub use extension::{DistExtension, ExtensionError, SourceDistExtension};
pub use source_dist::{SourceDistFilename, SourceDistFilenameError};
pub use wheel::{TagSet, WheelFilename, WheelFilenameError};
pub use wheel::{WheelFilename, WheelFilenameError};

mod build_tag;
mod egg;
mod extension;
mod source_dist;
mod splitter;
mod wheel;
mod wheel_tag;

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum DistFilename {
Expand Down Expand Up @@ -97,13 +98,9 @@ impl Display for DistFilename {
#[cfg(test)]
mod tests {
use crate::WheelFilename;
use uv_platform_tags::{AbiTag, LanguageTag, PlatformTag};

#[test]
fn wheel_filename_size() {
assert_eq!(size_of::<WheelFilename>(), 72);
assert_eq!(size_of::<LanguageTag>(), 16);
assert_eq!(size_of::<AbiTag>(), 16);
assert_eq!(size_of::<PlatformTag>(), 16);
assert_eq!(size_of::<WheelFilename>(), 48);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ Ok(
platform_tag: [
Any,
],
repr: "202206090410-py3-none-any",
},
},
},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ Ok(
arch: X86_64,
},
],
repr: "cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64",
},
},
},
Expand Down
173 changes: 25 additions & 148 deletions crates/uv-distribution-filename/src/wheel.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ use uv_platform_tags::{
};

use crate::splitter::MemchrSplitter;
use crate::wheel_tag::{WheelTag, WheelTagLarge, WheelTagSmall};
use crate::{BuildTag, BuildTagError};

#[derive(
Expand Down Expand Up @@ -177,7 +178,7 @@ impl WheelFilename {
));
};

let (name, version, build_tag, python_tag, abi_tag, platform_tag, is_large) =
let (name, version, build_tag, python_tag, abi_tag, platform_tag, is_small) =
if let Some(platform_tag) = splitter.next() {
if splitter.next().is_some() {
return Err(WheelFilenameError::InvalidWheelFileName(
Expand All @@ -193,7 +194,7 @@ impl WheelFilename {
&stem[abi_tag_or_platform_tag + 1..platform_tag],
&stem[platform_tag + 1..],
// Always take the slow path if a build tag is present.
true,
false,
)
} else {
(
Expand All @@ -206,7 +207,7 @@ impl WheelFilename {
// Determine whether any of the tag types contain a period, which would indicate
// that at least one of the tag types includes multiple tags (which in turn
// necessitates taking the slow path).
memchr(b'.', stem[build_tag_or_python_tag..].as_bytes()).is_some(),
memchr(b'.', stem[build_tag_or_python_tag..].as_bytes()).is_none(),
)
};

Expand All @@ -221,44 +222,38 @@ impl WheelFilename {
})
.transpose()?;

let tags = if is_large {
let tags = if let Some(small) = is_small
.then(|| {
Some(WheelTagSmall {
python_tag: LanguageTag::from_str(python_tag).ok()?,
abi_tag: AbiTag::from_str(abi_tag).ok()?,
platform_tag: PlatformTag::from_str(platform_tag).ok()?,
})
})
.flatten()
{
WheelTag::Small { small }
} else {
// Store the plaintext representation of the tags.
let repr = &stem[build_tag_or_python_tag + 1..];
WheelTag::Large {
large: Box::new(WheelTagLarge {
build_tag,
python_tag: MemchrSplitter::split(python_tag, b'.')
.map(LanguageTag::from_str)
.collect::<Result<_, _>>()
.map_err(|err| {
WheelFilenameError::InvalidLanguageTag(filename.to_string(), err)
})?,
.filter_map(Result::ok)
.collect(),
abi_tag: MemchrSplitter::split(abi_tag, b'.')
.map(AbiTag::from_str)
.collect::<Result<_, _>>()
.map_err(|err| {
WheelFilenameError::InvalidAbiTag(filename.to_string(), err)
})?,
.filter_map(Result::ok)
.collect(),
platform_tag: MemchrSplitter::split(platform_tag, b'.')
.map(PlatformTag::from_str)
.collect::<Result<_, _>>()
.map_err(|err| {
WheelFilenameError::InvalidPlatformTag(filename.to_string(), err)
})?,
.filter_map(Result::ok)
.collect(),
repr: repr.into(),
}),
}
} 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 {
Expand Down Expand Up @@ -311,124 +306,6 @@ 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 {
if let Some(build_tag) = &self.build_tag {
write!(f, "{build_tag}-")?;
}
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}")]
Expand Down
113 changes: 113 additions & 0 deletions crates/uv-distribution-filename/src/wheel_tag.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,113 @@
use std::fmt::{Display, Formatter};

use crate::BuildTag;
use uv_platform_tags::{AbiTag, LanguageTag, PlatformTag};
use uv_small_str::SmallString;

/// 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(crate) 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, a build
/// tag, or an unsupported tag (i.e., a tag that can't be represented by [`LanguageTag`],
/// [`AbiTag`], or [`PlatformTag`]). (Unsupported tags are filtered out, but retained in the display
/// representation of [`WheelTagLarge`].)
#[derive(
Debug,
Clone,
Eq,
PartialEq,
Ord,
PartialOrd,
Hash,
rkyv::Archive,
rkyv::Deserialize,
rkyv::Serialize,
)]
#[rkyv(derive(Debug))]
pub(crate) 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)]
pub(crate) struct WheelTagSmall {
/// The Python tag, e.g., `py3` in `1.2.3-py3-none-any`.
pub(crate) python_tag: LanguageTag,
/// The ABI tag, e.g., `none` in `1.2.3-py3-none-any`.
pub(crate) abi_tag: AbiTag,
/// The platform tag, e.g., `none` in `1.2.3-py3-none-any`.
pub(crate) 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(crate) struct WheelTagLarge {
/// The optional build tag, e.g., `73` in `1.2.3-73-py3-none-any-none`.
pub(crate) build_tag: Option<BuildTag>,
/// The Python tag(s), e.g., `py3` in `1.2.3-73-py3-none-any-none`.
pub(crate) python_tag: TagSet<LanguageTag>,
/// The ABI tag(s), e.g., `none` in `1.2.3-73-py3-none-any-none`.
pub(crate) abi_tag: TagSet<AbiTag>,
/// The platform tag(s), e.g., `none` in `1.2.3-73-py3-none-any-none`.
pub(crate) platform_tag: TagSet<PlatformTag>,
/// The string representation of the tag.
pub(crate) repr: SmallString,
}

impl Display for WheelTagLarge {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.repr)
}
}
4 changes: 2 additions & 2 deletions crates/uv-distribution-types/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1343,8 +1343,8 @@ mod test {
/// Ensure that we don't accidentally grow the `Dist` sizes.
#[test]
fn dist_size() {
assert!(size_of::<Dist>() <= 208, "{}", size_of::<Dist>());
assert!(size_of::<BuiltDist>() <= 208, "{}", size_of::<BuiltDist>());
assert!(size_of::<Dist>() <= 200, "{}", size_of::<Dist>());
assert!(size_of::<BuiltDist>() <= 200, "{}", size_of::<BuiltDist>());
assert!(
size_of::<SourceDist>() <= 176,
"{}",
Expand Down
Loading

0 comments on commit 7740d8d

Please sign in to comment.