Fixes for fix-wheel script (#97)

Author: cretz

scripts/fix_wheel.py

@@ -1,57 +1,66 @@
 import os
-import sys
-import zipfile
+import shutil
+import subprocess
 from glob import glob
 
 if __name__ == "__main__":
     # Due to https://github.com/python-poetry/poetry/issues/3509 and Poetry
-    # assuming the tags, we have to change the wheel ourselves after build
+    # assuming the tags, we have to change the wheel ourselves after build. In
+    # order to keep checksums proper, we use wheel pack/unpack.
 
     # Get the file from the dist dir
     dist_files = glob("dist/*.whl")
-    if not dist_files:
-        raise RuntimeError("Should have at least one wheel file in dist")
-    for existing_wheel_file in dist_files:
-        # Rename the wheel file and confirm it is changed. We need to make it
-        # py37 minimum interpreter and abi3 compat.
-        wheel_file_pieces = existing_wheel_file.split("-")
-        if len(wheel_file_pieces) < 4:
-            raise RuntimeError("Expecting at least 4 wheel pieces")
-        wheel_file_pieces[2] = "cp37"
-        wheel_file_pieces[3] = "abi3"
-        new_wheel_file = "-".join(wheel_file_pieces)
-        if existing_wheel_file == new_wheel_file:
-            raise RuntimeError("Wheel file already fixed")
-        print(
-            f"Converting from {existing_wheel_file} to {new_wheel_file}",
-            file=sys.stderr,
-        )
-
-        # Walk the Zip writing files except the WHEEL file which we must alter the
-        # tag on the WHEEL file
-        with zipfile.ZipFile(new_wheel_file, "w") as zipwrite, zipfile.ZipFile(
-            existing_wheel_file, "r"
-        ) as zipread:
-            found_wheel_tag = False
-            for item in zipread.infolist():
-                data = zipread.read(item.filename)
-                _, _, filename = item.filename.rpartition("/")
-                # Change the WHEEL tag
-                if filename == "WHEEL":
-                    lines = data.splitlines()
-                    for i in range(len(lines)):
-                        if lines[i].startswith(b"Tag: "):
-                            pieces = lines[i][len("Tag: ") :].split(b"-")
-                            if len(pieces) < 3:
-                                raise RuntimeError(
-                                    "Expecting at least 3 wheel tag pieces"
-                                )
-                            pieces[0] = b"cp37"
-                            pieces[1] = b"abi3"
-                            lines[i] = b"Tag: " + b"-".join(pieces)
-                            found_wheel_tag = True
-                    data = b"\n".join(lines)
-                zipwrite.writestr(item, data)
-            if not found_wheel_tag:
-                raise RuntimeError("Did not find WHEEL file with tag entry")
-        os.remove(existing_wheel_file)
+    if len(dist_files) != 1:
+        raise RuntimeError(f"Should have only one wheel file, found: {dist_files}")
+
+    # Run unpack into temp directory
+    if os.path.exists("dist/temp"):
+        raise RuntimeError("dist/temp directory already present")
+    subprocess.check_call(["wheel", "unpack", "--dest", "dist/temp", dist_files[0]])
+
+    # Read WHEEL contents
+    wheel_files = glob("dist/temp/*/*.dist-info/WHEEL")
+    if len(wheel_files) != 1:
+        raise RuntimeError(f"Should have only one WHEEL file, found: {wheel_files}")
+    with open(wheel_files[0]) as f:
+        wheel_lines = f.read().splitlines()
+
+    # Alter the "Tag" to use 3.7+ Python, ABI3 limited API, or if macOS ARM,
+    # manually set as 11_0 w/ 3.8+
+    found_wheel_tag = False
+    for i, line in enumerate(wheel_lines):
+        if line.startswith("Tag: "):
+            pieces = line[len("Tag: ") :].split("-")
+            if len(pieces) < 3:
+                raise RuntimeError("Expecting at least 3 wheel tag pieces")
+            pieces[1] = "abi3"
+            if pieces[2].startswith("macosx_") and pieces[2].endswith("_arm64"):
+                pieces[0] = "cp38"
+                pieces[2] = "macosx_11_0_arm64"
+            else:
+                pieces[0] = "cp37"
+            wheel_lines[i] = "Tag: " + "-".join(pieces)
+            found_wheel_tag = True
+            break
+    if not found_wheel_tag:
+        raise RuntimeError("Could not find WHEEL tag")
+
+    # Write the WHEEL file
+    with open(wheel_files[0], "w") as f:
+        f.write("\n".join(wheel_lines))
+
+    # Pack the wheel
+    unpacked_dirs = glob("dist/temp/*")
+    subprocess.check_call(["wheel", "pack", "--dest", "dist", unpacked_dirs[0]])
+
+    # Remove temp dir
+    shutil.rmtree("dist/temp")
+
+    # If there are multiple wheels now, remove the old one
+    new_dist_files = glob("dist/*.whl")
+    new_dist_files.remove(dist_files[0])
+    if new_dist_files:
+        os.remove(dist_files[0])
+        print(f"Created wheel {new_dist_files[0]} from {dist_files[0]}")
+    else:
+        print(f"Overwrote wheel {dist_files[0]}")