Merge pull request #12 from pelson/feature/generalise_wheel_update

Generalise the wheel update procedure, thus making it easy for other backends to do the same tricks

.github/workflows/tests.yml

@@ -13,9 +13,16 @@ jobs:
     runs-on: ${{ matrix.os }}
 
     strategy:
+      fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
-        python-version: ["3.7", "3.12"]
+        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"]
+        exclude:
+          - os: macos-latest
+            python-version: "3.7"
+        include:
+          - os: macos-13
+            python-version: "3.7"
 
     steps:
       - uses: "actions/checkout@v4"

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "setuptools-ext"
-version = "0.6"
+version = "0.7"
 requires-python = ">= 3.7"
 description = "Extension of setuptools to support all core metadata fields"
 authors = [

setuptools_ext.py

@@ -1,17 +1,21 @@
 """Extension of setuptools to support all core metadata fields"""
+
 import base64
 import email.policy
 import hashlib
+import shutil
+import sys
+import typing
+import zipfile
 from pathlib import Path
-from zipfile import ZipFile
 
-from setuptools.build_meta import *
+from setuptools.build_meta import *  # noqa
 from setuptools.build_meta import build_wheel as orig_build_wheel
-try:
-    # stdlib Python 3.11+
-    import tomllib as toml
-except ImportError:
-    import tomli as toml
+
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
 
 
 allowed_fields = {
@@ -28,7 +32,7 @@
 
 
 def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
-    project = toml.loads(Path("pyproject.toml").read_text())
+    project = tomllib.loads(Path("pyproject.toml").read_text())
     ours = project.get("tool", {}).get("setuptools-ext", {})
     extra_metadata = {}
     for key, vals in ours.items():
@@ -49,6 +53,9 @@ def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
 
 
 def rewrite_metadata(data, extra_metadata):
+    """
+    Rewrite the METADATA file to include the given additional metadata.
+    """
     pkginfo = email.message_from_string(data.decode())
     # delete some annoying kv that distutils seems to put in there for no reason
     for key in dict(pkginfo):
@@ -66,14 +73,121 @@ def rewrite_metadata(data, extra_metadata):
     return result
 
 
-def rewrite_record(data, new_line):
-    lines = []
-    for line in data.decode().splitlines():
-        fname = line.split(",")[0]
-        if fname.endswith(".dist-info/METADATA"):
-            line = new_line
-        lines.append(line)
-    return "\n".join(lines).encode()
+class WheelRecord:
+    """
+    Represents the RECORD file of a wheel, which can be updated with new checksums
+    using the record_file method.
+
+    See also the (limited) spec on RECORD at
+    https://packaging.python.org/en/latest/specifications/binary-distribution-format/#signed-wheel-files
+    """
+
+    def __init__(self, record_content: str = ""):
+        #: Records mapping filename to (hash, length) tuples.
+        self._records: typing.Dict[str, typing.Tuple[str, str]] = {}
+        self.update_from_record(record_content)
+
+    def update_from_record(
+        self, record_content: typing.Union[str, "WheelRecord"]
+    ) -> None:
+        """
+        Update this WheelRecord given another WheelRecord, or RECORD contents
+        """
+        if isinstance(record_content, WheelRecord):
+            record_content = record_content.record_contents()
+        for line in record_content.splitlines():
+            path, file_hash, length = line.split(",")
+            self._records[path] = file_hash, length
+
+    def record_file(self, filename, file_content: typing.Union[bytes, str]):
+        """
+        Record the filename and appropriate digests of its contents
+        """
+        if isinstance(file_content, str):
+            file_content = file_content.encode("utf-8")
+        digest = hashlib.sha256(file_content).digest()
+        checksum = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
+        self._records[filename] = f"sha256={checksum}", str(len(file_content))
+
+    def record_contents(self) -> str:
+        """
+        Output the representation of this WheelRecord as a string, suitable for
+        writing to a RECORD file.
+        """
+        contents = []
+        for path, (file_hash, length) in self._records.items():
+            contents.append(f"{path},{file_hash},{length}")
+        return "\n".join(contents)
+
+
+class WheelModifier:
+    """
+    Representation of an existing wheel with lazily modified contents that
+    can be written on-demand with the write_wheel method.
+    """
+
+    def __init__(self, wheel_zipfile: zipfile.ZipFile):
+        self._wheel_zipfile = wheel_zipfile
+        # Track updated file contents.
+        self._updates: typing.Dict[str, typing.Tuple[zipfile.ZipInfo, bytes]] = {}
+
+    def dist_info_dirname(self):
+        for filename in self._wheel_zipfile.namelist():
+            # TODO: We could use the filename of the zipfile... but we don't
+            #  necessarily have it.
+            if filename.endswith(".dist-info/METADATA"):
+                return filename.rsplit("/", 1)[0]
+
+    def read(self, filename: str) -> bytes:
+        if filename in self._updates:
+            return self._updates[filename][1]
+        else:
+            return self._wheel_zipfile.read(filename)
+
+    def zipinfo(self, filename: str) -> zipfile.ZipInfo:
+        if filename in self._updates:
+            return self._updates[filename][0]
+        return self._wheel_zipfile.getinfo(filename)
+
+    def write(
+        self,
+        filename: typing.Union[str, zipfile.ZipInfo],
+        content: bytes,
+    ) -> None:
+        zinfo: zipfile.ZipInfo
+        if isinstance(filename, zipfile.ZipInfo):
+            zinfo = filename
+        else:
+            try:
+                zinfo = self.zipinfo(filename)
+            except KeyError:
+                raise ValueError(
+                    f"Unable to write filename {filename} as there is no existing "
+                    "file information in the archive. Please provide a zipinfo "
+                    "instance when writing."
+                )
+        self._updates[typing.cast(str, zinfo.filename)] = zinfo, content
+
+    def write_wheel(self, file: typing.Union[str, Path, typing.IO[bytes]]) -> None:
+        distinfo_dir = self.dist_info_dirname()
+        record_filename = f"{distinfo_dir}/RECORD"
+        orig_record = WheelRecord(self.read(record_filename).decode())
+        with zipfile.ZipFile(file, "w") as z_out:
+            for zinfo in self._wheel_zipfile.infolist():
+                if zinfo.filename == record_filename:
+                    # We deal with record last.
+                    continue
+                if zinfo.filename in self._updates:
+                    zinfo, content = self._updates.pop(zinfo.filename)
+                    orig_record.record_file(zinfo.filename, content)
+                else:
+                    content = self._wheel_zipfile.read(zinfo.filename)
+                z_out.writestr(zinfo, content)
+            for zinfo, content in self._updates.values():
+                orig_record.record_file(zinfo.filename, content)
+                z_out.writestr(zinfo, content)
+            record_zinfo = self._wheel_zipfile.getinfo(record_filename)
+            z_out.writestr(record_zinfo, orig_record.record_contents())
 
 
 def rewrite_whl(path, extra_metadata):
@@ -87,24 +201,13 @@ def rewrite_whl(path, extra_metadata):
     # be changed later on, but unfortunately for now the only option is to rewrite the
     # generated .whl with our modifications
     tmppath = path.parent.joinpath("." + path.name)
-    checksum = record = None
-    with ZipFile(str(path), "r") as z_in, ZipFile(str(tmppath), "w") as z_out:
-        z_out.comment = z_in.comment
-        for zinfo in z_in.infolist():
-            data = z_in.read(zinfo.filename)
-            if zinfo.filename.endswith(".dist-info/METADATA"):
-                data = rewrite_metadata(data, extra_metadata)
-                digest = hashlib.sha256(data).digest()
-                checksum = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
-                new_line = f"{zinfo.filename},sha256={checksum},{len(data)}"
-            if zinfo.filename.endswith(".dist-info/RECORD"):
-                record = zinfo, data
-                continue
-            z_out.writestr(zinfo, data)
-        if record is not None:
-            record_info, record_data = record
-            if checksum is not None:
-                record_data = rewrite_record(record_data, new_line)
-            z_out.writestr(record_info, record_data)
-    path.write_bytes(tmppath.read_bytes())
-    tmppath.unlink()
+
+    with zipfile.ZipFile(str(path), "r") as whl_zip:
+        whl = WheelModifier(whl_zip)
+        metadata_filename = f"{whl.dist_info_dirname()}/METADATA"
+        metadata = rewrite_metadata(whl.read(metadata_filename), extra_metadata)
+        whl.write(metadata_filename, metadata)
+        with tmppath.open("wb") as whl_fh:
+            whl.write_wheel(whl_fh)
+
+    shutil.move(tmppath, path)