From 58ada5c660571a0605f15c1f912c0a103f050813 Mon Sep 17 00:00:00 2001 From: pmattsso Date: Wed, 8 Jul 2026 14:15:08 +0100 Subject: [PATCH] fix(downloads): accept wheels whose dist-info omits the local version segment Some vendor-built wheels (e.g. AMD's tensorflow-rocm) name their dist-info directory using only the base version while the wheel filename includes a PEP 440 local segment. wheelfile.WheelFile rejects these with "Missing RECORD file" because it expects the directory name to match the full filename version. Add a fallback that catches WheelError and checks for RECORD and METADATA under the base-version dist-info directory when the filename carries a local version. The wheel is accepted with a warning if both files exist. Closes: #1239 Signed-off-by: pmattsso Co-authored-by: Claude Opus 4.6 --- src/fromager/downloads.py | 68 +++++++++++++++++++++++++++++++-------- tests/test_downloads.py | 49 ++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 13 deletions(-) diff --git a/src/fromager/downloads.py b/src/fromager/downloads.py index 82dd7b8d..3f9dbdca 100644 --- a/src/fromager/downloads.py +++ b/src/fromager/downloads.py @@ -144,6 +144,41 @@ def download_sdist( return filepath +def _validate_wheel_with_local_version(filepath: pathlib.Path, basename: str) -> bool: + """Check whether a wheel with a local version has valid dist-info under the base version. + + Some third-party wheels name their dist-info directory using only the + public (base) version (e.g. ``pkg-1.0.dist-info``) while the wheel + filename carries a local segment (e.g. ``pkg-1.0+local``). + ``wheelfile.WheelFile`` rejects these because it expects the directory + name to match the full filename version. + + Returns ``True`` when the wheel has a local version **and** the + base-version dist-info directory contains both ``RECORD`` and + ``METADATA``; ``False`` otherwise. + """ + _, version, _, _ = parse_wheel_filename(basename) + if not version.local: + return False + + dist_name = basename.split("-", 1)[0] + base_dist_info = f"{dist_name}-{version.public}.dist-info" + with zipfile.ZipFile(filepath) as zf: + names = zf.namelist() + record_path = f"{base_dist_info}/RECORD" + metadata_path = f"{base_dist_info}/METADATA" + if record_path in names and metadata_path in names: + logger.warning( + "wheel %s has dist-info directory %s instead of expected %s-%s.dist-info", + basename, + base_dist_info, + dist_name, + version, + ) + return True + return False + + def download_wheel( *, destination_dir: pathlib.Path, @@ -152,16 +187,16 @@ def download_wheel( ) -> pathlib.Path: """Download a wheel and verify it. - Validates that the filename is a valid wheel name (using - :func:`packaging.utils.parse_wheel_filename`) before downloading, - then checks that the zip archive is non-empty. - - .. note:: + Validates the filename with + :func:`packaging.utils.parse_wheel_filename`, downloads the file, + then opens it with :class:`wheel.wheelfile.WheelFile` to verify + that the dist-info directory and ``RECORD`` / ``METADATA`` files + exist. - The function does not validate dist-info ``METADATA`` and other - files described in the `binary distribution format - `_ - specification. + When the wheel carries a local version segment in its filename but + the dist-info directory uses only the base version (a pattern seen + in some third-party builds), the strict ``WheelFile`` check is + relaxed via :func:`_validate_wheel_with_local_version`. .. versionadded:: 0.90.0 """ @@ -182,10 +217,17 @@ def download_wheel( # NOTE: Does not validate that METADATA file is valid or consistent with # wheel's distribution name and version. # https://packaging.python.org/en/latest/specifications/binary-distribution-format/ - with wheelfile.WheelFile(filepath) as wf: - di_metadata = f"{wf.dist_info_path}/METADATA" - if di_metadata not in wf.namelist(): - raise wheelfile.WheelError(f"{di_metadata!r} missing") + try: + with wheelfile.WheelFile(filepath) as wf: + di_metadata = f"{wf.dist_info_path}/METADATA" + if di_metadata not in wf.namelist(): + raise wheelfile.WheelError(f"{di_metadata!r} missing") + except wheelfile.WheelError: + # Some wheels omit the local version segment from the dist-info + # directory name (e.g. "pkg-1.0.dist-info" when the filename says + # "pkg-1.0+local"). Fall back to a manual check. + if not _validate_wheel_with_local_version(filepath, basename): + raise return filepath diff --git a/tests/test_downloads.py b/tests/test_downloads.py index 27ad1d97..a55335b0 100644 --- a/tests/test_downloads.py +++ b/tests/test_downloads.py @@ -13,6 +13,7 @@ from wheel.wheelfile import WheelFile from fromager.downloads import ( + _validate_wheel_with_local_version, download_git_source, download_sdist, download_url, @@ -201,6 +202,54 @@ def test_download_wheel_missing_metadata(tmp_path: pathlib.Path) -> None: ) +def test_download_wheel_local_version_mismatch(tmp_path: pathlib.Path) -> None: + """Wheel whose dist-info uses base version while filename has local.""" + filename = "pkg-1.0+local123-py3-none-any.whl" + f = tmp_path / filename + with zipfile.ZipFile(f, "w") as zf: + zf.writestr( + "pkg-1.0.dist-info/METADATA", + "Metadata-Version: 1.0\nName: pkg\nVersion: 1.0+local123\n", + ) + zf.writestr("pkg-1.0.dist-info/WHEEL", "Wheel-Version: 1.0\n") + zf.writestr("pkg-1.0.dist-info/RECORD", "") + with _patch_download(f): + result = download_wheel( + destination_dir=tmp_path, + url=f"https://pkg.test/{filename}", + ) + assert result == f + + +def test_validate_wheel_with_local_version_valid(tmp_path: pathlib.Path) -> None: + """Fallback validation passes for a well-formed wheel with base-only dist-info.""" + filename = "pkg-1.0+local-py3-none-any.whl" + f = tmp_path / filename + with zipfile.ZipFile(f, "w") as zf: + zf.writestr("pkg-1.0.dist-info/METADATA", "Metadata-Version: 1.0\n") + zf.writestr("pkg-1.0.dist-info/RECORD", "") + assert _validate_wheel_with_local_version(f, filename) is True + + +def test_validate_wheel_with_local_version_no_local(tmp_path: pathlib.Path) -> None: + """Fallback returns False when the filename has no local version.""" + filename = "pkg-1.0-py3-none-any.whl" + f = tmp_path / filename + with zipfile.ZipFile(f, "w") as zf: + zf.writestr("pkg-1.0.dist-info/METADATA", "Metadata-Version: 1.0\n") + zf.writestr("pkg-1.0.dist-info/RECORD", "") + assert _validate_wheel_with_local_version(f, filename) is False + + +def test_validate_wheel_with_local_version_no_record(tmp_path: pathlib.Path) -> None: + """Fallback returns False when RECORD is missing from base-version dist-info.""" + filename = "pkg-1.0+local-py3-none-any.whl" + f = tmp_path / filename + with zipfile.ZipFile(f, "w") as zf: + zf.writestr("pkg-1.0.dist-info/METADATA", "Metadata-Version: 1.0\n") + assert _validate_wheel_with_local_version(f, filename) is False + + # -- download_git_source ------------------------------------------------------ _MOCK_GIT_CLONE = "fromager.downloads.gitutils.git_clone_fast"