fix: setup.py --all only put local driver if arch matches (microsoft#1043)

mxschmitt · web-flow · commit 37c0533e04e3 · 2021-11-19T23:36:11.000+01:00
diff --git a/setup.py b/setup.py
@@ -14,11 +14,13 @@
 
 import glob
 import os
+import platform
 import shutil
 import subprocess
 import sys
 import zipfile
 from pathlib import Path
+from typing import Dict, List
 
 from setuptools import find_packages, setup
 
@@ -40,6 +42,23 @@ def extractall(zip: zipfile.ZipFile, path: str) -> None:
             os.chmod(extracted_path, attr)
 
 
+def download_driver(zip_name: str) -> None:
+    zip_file = f"playwright-{driver_version}-{zip_name}.zip"
+    if os.path.exists("driver/" + zip_file):
+        return
+    url = "https://playwright.azureedge.net/builds/driver/"
+    if (
+        "-alpha" in driver_version
+        or "-beta" in driver_version
+        or "-next" in driver_version
+    ):
+        url = url + "next/"
+    url = url + zip_file
+    print(f"Fetching {url}")
+    # Don't replace this with urllib - Python won't have certificates to do SSL on all platforms.
+    subprocess.check_call(["curl", url, "-o", "driver/" + zip_file])
+
+
 class PlaywrightBDistWheelCommand(BDistWheelCommand):
     user_options = BDistWheelCommand.user_options + [
         ("all", "a", "create wheels for all platforms")
@@ -57,97 +76,90 @@ def run(self) -> None:
         super().run()
         os.makedirs("driver", exist_ok=True)
         os.makedirs("playwright/driver", exist_ok=True)
-        if self.all:
-            # If building for all platforms
-            platform_map = {
-                "darwin": [
-                    {
-                        "zip_name": "mac",
-                        "wheels": [
-                            "macosx_10_13_x86_64.whl",
-                            "macosx_11_0_universal2.whl",
-                        ],
-                    },
-                    {
-                        "zip_name": "mac-arm64",
-                        "wheels": [
-                            "macosx_11_0_arm64.whl",
-                        ],
-                    },
-                ],
-                "linux": [
-                    {"zip_name": "linux", "wheels": ["manylinux1_x86_64.whl"]},
-                    {
-                        "zip_name": "linux-arm64",
-                        "wheels": ["manylinux_2_17_aarch64.manylinux2014_aarch64.whl"],
-                    },
-                ],
-                "win32": [
-                    {
-                        "zip_name": "win32_x64",
-                        "wheels": ["win32.whl", "win_amd64.whl"],
-                    }
-                ],
-            }
-            platforms = [*platform_map.values()]
-        else:
-            # If building for only current platform
-            platform_map = {
-                "darwin": [
-                    {
-                        "zip_name": "mac",
-                        "wheels": ["macosx_10_13_x86_64.whl"],
-                    },
-                ],
-                "linux": [{"zip_name": "linux", "wheels": ["manylinux1_x86_64.whl"]}],
-                "win32": [
-                    {
-                        "zip_name": "win32_x64",
-                        "wheels": ["win_amd64.whl"],
-                    }
-                ],
-            }
-            platforms = [platform_map[sys.platform]]
-        for platform in platforms:
-            for arch in platform:
-                zip_file = f"playwright-{driver_version}-{arch['zip_name']}.zip"
-                if not os.path.exists("driver/" + zip_file):
-                    url = "https://playwright.azureedge.net/builds/driver/"
-                    if (
-                        "-alpha" in driver_version
-                        or "-beta" in driver_version
-                        or "-next" in driver_version
-                    ):
-                        url = url + "next/"
-                    url = url + zip_file
-                    print(f"Fetching {url}")
-                    # Don't replace this with urllib - Python won't have certificates to do SSL on all platforms.
-                    subprocess.check_call(["curl", url, "-o", "driver/" + zip_file])
-        base_wheel_location = glob.glob(os.path.join(self.dist_dir, "*.whl"))[0]
-        without_platform = base_wheel_location[:-7]
+        base_wheel_bundles: List[Dict[str, str]] = [
+            {
+                "wheel": "macosx_10_13_x86_64.whl",
+                "machine": "x86_64",
+                "platform": "darwin",
+                "zip_name": "mac",
+            },
+            {
+                "wheel": "macosx_11_0_universal2.whl",
+                "machine": "x86_64",
+                "platform": "darwin",
+                "zip_name": "mac",
+            },
+            {
+                "wheel": "macosx_11_0_arm64.whl",
+                "machine": "arm64",
+                "platform": "darwin",
+                "zip_name": "mac-arm64",
+            },
+            {
+                "wheel": "manylinux1_x86_64.whl",
+                "machine": "x86_64",
+                "platform": "linux",
+                "zip_name": "linux",
+            },
+            {
+                "wheel": "manylinux_2_17_aarch64.manylinux2014_aarch64.whl",
+                "machine": "aarch64",
+                "platform": "linux",
+                "zip_name": "linux-arm64",
+            },
+            {
+                "wheel": "win32.whl",
+                "machine": "i386",
+                "platform": "win32",
+                "zip_name": "win32_x64",
+            },
+            {
+                "wheel": "win_amd64.whl",
+                "machine": "amd64",
+                "platform": "win32",
+                "zip_name": "win32_x64",
+            },
+        ]
+        self._download_and_extract_local_driver(base_wheel_bundles)
 
-        for platform in platforms:
-            for arch in platform:
-                zip_file = f"driver/playwright-{driver_version}-{arch['zip_name']}.zip"
-                with zipfile.ZipFile(zip_file, "r") as zip:
-                    extractall(zip, f"driver/{arch['zip_name']}")
-                if platform_map[sys.platform][0] in platform:
-                    with zipfile.ZipFile(zip_file, "r") as zip:
-                        extractall(zip, "playwright/driver")
-                for wheel in arch["wheels"]:
-                    wheel_location = without_platform + wheel
-                    shutil.copy(base_wheel_location, wheel_location)
-                    with zipfile.ZipFile(wheel_location, "a") as zip:
-                        driver_root = os.path.abspath(f"driver/{arch['zip_name']}")
-                        for dir_path, _, files in os.walk(driver_root):
-                            for file in files:
-                                from_path = os.path.join(dir_path, file)
-                                to_path = os.path.relpath(from_path, driver_root)
-                                zip.write(from_path, f"playwright/driver/{to_path}")
-                        zip.writestr(
-                            "playwright/driver/README.md", f"{wheel} driver package"
-                        )
+        wheels = base_wheel_bundles
+        if not self.all:
+            # Limit to 1, since for MacOS e.g. we have multiple wheels for the same platform and architecture and Conda expects 1.
+            wheels = list(
+                filter(
+                    lambda wheel: wheel["platform"] == sys.platform
+                    and wheel["machine"] == platform.machine().lower(),
+                    base_wheel_bundles,
+                )
+            )[:1]
+        self._build_wheels(wheels)
 
+    def _build_wheels(
+        self,
+        wheels: List[Dict[str, str]],
+    ) -> None:
+        base_wheel_location: str = glob.glob(os.path.join(self.dist_dir, "*.whl"))[0]
+        without_platform = base_wheel_location[:-7]
+        for wheel_bundle in wheels:
+            download_driver(wheel_bundle["zip_name"])
+            zip_file = (
+                f"driver/playwright-{driver_version}-{wheel_bundle['zip_name']}.zip"
+            )
+            with zipfile.ZipFile(zip_file, "r") as zip:
+                extractall(zip, f"driver/{wheel_bundle['zip_name']}")
+            wheel_location = without_platform + wheel_bundle["wheel"]
+            shutil.copy(base_wheel_location, wheel_location)
+            with zipfile.ZipFile(wheel_location, "a") as zip:
+                driver_root = os.path.abspath(f"driver/{wheel_bundle['zip_name']}")
+                for dir_path, _, files in os.walk(driver_root):
+                    for file in files:
+                        from_path = os.path.join(dir_path, file)
+                        to_path = os.path.relpath(from_path, driver_root)
+                        zip.write(from_path, f"playwright/driver/{to_path}")
+                zip.writestr(
+                    "playwright/driver/README.md",
+                    f"{wheel_bundle['wheel']} driver package",
+                )
         os.remove(base_wheel_location)
         if InWheel:
             for whlfile in glob.glob(os.path.join(self.dist_dir, "*.whl")):
@@ -163,6 +175,27 @@ def run(self) -> None:
         else:
             print("auditwheel not installed, not updating RECORD file")
 
+    def _download_and_extract_local_driver(
+        self,
+        wheels: List[Dict[str, str]],
+    ) -> None:
+        zip_names_for_current_system = set(
+            map(
+                lambda wheel: wheel["zip_name"],
+                filter(
+                    lambda wheel: wheel["machine"] == platform.machine().lower()
+                    and wheel["platform"] == sys.platform,
+                    wheels,
+                ),
+            )
+        )
+        assert len(zip_names_for_current_system) == 1
+        zip_name = zip_names_for_current_system.pop()
+        download_driver(zip_name)
+        zip_file = f"driver/playwright-{driver_version}-{zip_name}.zip"
+        with zipfile.ZipFile(zip_file, "r") as zip:
+            extractall(zip, "playwright/driver")
+
 
 setup(
     name="playwright",
