Tidy up script.py and import_tests.py

Author: hoodmane

packages/import_tests.py

@@ -1,27 +1,23 @@
 from pathlib import Path
 
-from typing import List
-import yaml
+from pyodide_build.io import MetaConfig
 
-def gen(packages: List[str], packages_dir = Path("packages")):
+
+def gen(packages: list[str], packages_dir: Path = Path("packages")) -> dict[str, list[str]]:
     res = {}
     for package in packages:
-        if package == "test": continue
+        if package == "test":
+            continue
         try:
-            with open(packages_dir / package / "meta.yaml", "r") as file:
-                imports = set()
-                meta = yaml.load(file, Loader=yaml.FullLoader)
-                # add package -> top-level if it exists
-                if "package" in meta:
-                    if "top-level" in meta["package"]:
-                        imports.update(meta["package"]["top-level"])
-                # add test -> imports if it exists
-                if "test" in meta:
-                    if "imports" in meta["test"]:
-                        imports.update(meta["test"]["imports"])
-                
-                if len(imports) > 0:
-                    res[package] = sorted(list(imports))
+            meta = MetaConfig.from_yaml(packages_dir / package / "meta.yaml")
         except FileNotFoundError:
             print(f"package {package}'s meta.yaml not found")
-    return res
+
+        imports = set()
+        # add package -> top-level if it exists
+        imports.update(meta.package.top_level)
+        imports.update(meta.test.imports)
+        if imports:
+            res[package] = sorted(list(imports))
+
+    return res

packages/script.py

@@ -1,29 +1,38 @@
-import os
+import hashlib
 import json
-from pathlib import Path
+import os
+import re
+import subprocess
+import sys
+import tarfile
 import tempfile
 import zipfile
-import tarfile
-import uuid
+from pathlib import Path
+from textwrap import dedent
+
 import boto3
 import re
 import sys
 import hashlib
 import time
 import requests
-from datetime import datetime
 
 import import_tests
 
-def normalize(name):
+
+def normalize(name: str) -> str:
     return re.sub(r"[-_.]+", "-", name).lower()
 
+
 # See setup.sh
 # prerequisite: emsdk, pyodide, packages -> pyodide/packages
 
-def gen_bzl_config(tag, dist):
-    bucket_url = "https://pyodide.edgeworker.net/python-package-bucket/" + tag + "/"
-    github_url = "https://github.com/cloudflare/pyodide-build-scripts/releases/download/" + tag + "/"
+
+def gen_bzl_config(tag: str, dist: Path) -> None:
+    bucket_url = f"https://pyodide.edgeworker.net/python-package-bucket/{tag}/"
+    github_url = (
+        f"https://github.com/cloudflare/pyodide-build-scripts/releases/download/{tag}/"
+    )
     lock_bytes = (dist / "pyodide-lock.json").read_bytes()
     lock_hash = hashlib.sha256(lock_bytes).hexdigest()
     zip_bytes = (dist / "pyodide_packages.tar.zip").read_bytes()
@@ -33,24 +42,36 @@ def gen_bzl_config(tag, dist):
 
     with open(dist / "pyodide-lock.json", "r") as file:
         lock = json.load(file)
-        packages = [package["name"] for package in lock["packages"].values()]
-        imports_to_test = import_tests.gen(packages)
-
-    with open("pyodide_bucket.bzl", "w") as f:
-        f.write("# Do not edit this file by hand. See docs/pyodide.md for info on how to generate it.\n")
-        f.write("# These variables are factored out here because they are being shared by the WORKSPACE files in\n")
-        f.write("# both edgeworker and workerd, as well as src/pyodide/BUILD.bazel\n")
-        f.write("PYODIDE_PACKAGE_BUCKET_URL = \"" + bucket_url + "\"\n")
-        f.write("PYODIDE_GITHUB_RELEASE_URL = \"" + github_url + "\"\n")
-        f.write("PYODIDE_LOCK_SHA256 = \"" + lock_hash + "\"\n")
-        f.write("PYODIDE_PACKAGES_TAR_ZIP_SHA256 = \"" + zip_hash + "\"\n")
-        f.write("PYODIDE_ALL_WHEELS_ZIP_SHA256 = \"" + all_wheels_hash + "\"\n\n")
-        f.write("# IMPORTANT: when updating this file in git, check the diff to make sure none of the imports below are being removed unexpectedly\n")
-        f.write("PYODIDE_IMPORTS_TO_TEST = " + json.dumps(imports_to_test, indent=3, sort_keys=True) + "\n")
+    packages = [package["name"] for package in lock["packages"].values()]
+    imports_to_test = import_tests.gen(packages)
+
+    Path("pyodide_bucket.bzl").write_text(
+        dedent(
+            f"""
+            # Do not edit this file by hand. See docs/pyodide.md for info on how to generate it.
+            # These variables are factored out here because they are being shared by the WORKSPACE files in
+            # both edgeworker and workerd, as well as src/pyodide/BUILD.bazel
+
+            PYODIDE_PACKAGE_BUCKET_URL = "{bucket_url}"
+            PYODIDE_GITHUB_RELEASE_URL = "{github_url}"
+            PYODIDE_LOCK_SHA256 = "{lock_hash}"
+            PYODIDE_PACKAGES_TAR_ZIP_SHA256 = "{zip_hash}"
+            PYODIDE_ALL_WHEELS_ZIP_SHA256 = "{all_wheels_hash}"
+
+            # IMPORTANT: when updating this file in git, check the diff to make sure none of the imports below are being removed unexpectedly
+
+            PYODIDE_IMPORTS_TO_TEST =
+            """
+        ).strip()
+        + " "
+        + json.dumps(imports_to_test, indent=3, sort_keys=True)
+        + "\n"
+    )
+
 
 # creates a package bundle .tar.zip file to be bundled in with edgeworker
 # the resulting bundle is written to dist/pyodide_packages.tar.zip
-def make_bundle(tag, dist = Path("dist")):
+def make_bundle(tag: str, dist: Path = Path("dist")) -> None:
     with open(dist / "pyodide-lock.json", "r") as file:
         lock = json.load(file)
     with tempfile.TemporaryDirectory(delete=False) as t:
@@ -62,8 +83,8 @@ def make_bundle(tag, dist = Path("dist")):
         for package in lock["packages"].values():
             name = normalize(package["name"])
             print("untarring " + name)
-            os.mkdir(tempdir / name)
-            if name.endswith("-tests") or name == "test": 
+            (tempdir / name).mkdir()
+            if name.endswith("-tests") or name == "test":
                 continue
             file = dist / package["file_name"]
             with tarfile.open(file, "r:gz") as zip:
@@ -72,27 +93,37 @@ def make_bundle(tag, dist = Path("dist")):
         with tarfile.open(tempdir / "pyodide_packages.tar", "w") as tar:
             tar.add(tempdir, arcname="./")
         # create zip file in dist/ from tarfile
-        with zipfile.ZipFile(dist / "pyodide_packages.tar.zip", "w", compression=zipfile.ZIP_DEFLATED) as zip:
+        with zipfile.ZipFile(
+            dist / "pyodide_packages.tar.zip", "w", compression=zipfile.ZIP_DEFLATED
+        ) as zip:
             zip.write(tempdir / "pyodide_packages.tar", "pyodide_packages.tar")
         # create all_wheels.zip file for testing
-        with zipfile.ZipFile(dist / "all_wheels.zip", "w", compression=zipfile.ZIP_DEFLATED) as zip:
+        with zipfile.ZipFile(
+            dist / "all_wheels.zip", "w", compression=zipfile.ZIP_DEFLATED
+        ) as zip:
             for package in lock["packages"].values():
                 file = dist / package["file_name"]
                 zip.write(file, f"{package['file_name']}")
-    
+
     gen_bzl_config(tag, dist)
 
+
 # uploads everything in dist to python-package-bucket at tag/...
-def upload_to_r2(tag, dist = Path("dist")):
+def upload_to_r2(tag: str, dist=Path("dist")) -> None:
     # upload to r2
-    s3 = boto3.client("s3", 
-                      endpoint_url = "https://" + os.environ.get("R2_ACCOUNT_ID") + ".r2.cloudflarestorage.com",
-                      aws_access_key_id = os.environ.get("R2_ACCESS_KEY_ID"),
-                      aws_secret_access_key = os.environ.get("R2_SECRET_ACCESS_KEY"),
-                      region_name="auto")
-    
+    r2_account_id = os.environ["R2_ACCOUNT_ID"]
+    r2_access_key = os.environ.get("R2_ACCESS_KEY_ID")
+    r2_secret_access_key = os.environ.get("R2_SECRET_ACCESS_KEY")
+    s3 = boto3.client(
+        "s3",
+        endpoint_url=f"https://{r2_account_id}.r2.cloudflarestorage.com",
+        aws_access_key_id=r2_access_key,
+        aws_secret_access_key=r2_secret_access_key,
+        region_name="auto",
+    )
+
     files_remaining = []
-    
+
     # upload entire dist directory to r2, excluding all_wheels.zip and pyodide_packages.tar.zip
     for root, dirs, files in os.walk(dist):
         for file in files:
@@ -101,30 +132,30 @@ def upload_to_r2(tag, dist = Path("dist")):
             path = Path(root) / file
             key = tag + "/" + str(path.relative_to(dist))
             files_remaining.append((path, key))
-    
+
     # attempt to upload each file 5 times. If after 5 attempts the file is still not accessible at pyodide.edgeworker.net then give up
     ATTEMPTS = 5
     for i in range(ATTEMPTS):
-        for (path, key) in files_remaining:
+        for path, key in files_remaining:
             print(f"uploading {path} to {key}")
             s3.upload_file(str(path), "python-package-bucket", key)
 
         new_files_remaining = []
 
         time.sleep(10)
 
-        for (path, key) in files_remaining:
+        for path, key in files_remaining:
             # Construct URL to fetch the uploaded file
             url = f"https://pyodide.edgeworker.net/python-package-bucket/{key}"
             print(f"Checking {url}")
-            
+
             try:
                 # Download the file content from the URL
                 response = requests.get(url)
                 response.raise_for_status()  # Raise an exception if the status is not 200 OK
-                
+
                 # Read the local file content
-                with open(path, 'rb') as f:
+                with open(path, "rb") as f:
                     local_content = f.read()
 
                 # Compare contents
@@ -143,19 +174,19 @@ def upload_to_r2(tag, dist = Path("dist")):
             break
 
         if i != ATTEMPTS - 1:
-            for (path, key) in files_remaining:
+            for path, key in files_remaining:
                 s3.delete_object(Bucket="python-package-bucket", Key=key)
 
     if files_remaining:
         raise Exception("Failed to upload packages after 5 attempts: ", files_remaining)
 
+
 # converts all the .zip wheels into .tar.gz format (destructively)
-def convert_wheels_to_tar_gz(dist = Path("dist")):
+def convert_wheels_to_tar_gz(dist: Path = Path("dist")) -> None:
     with open(dist / "pyodide-lock.json", "r") as file:
         lock = json.load(file)
-    
+
     for package in lock["packages"].values():
-        name = normalize(package["name"])
         file = dist / package["file_name"]
         # check file ends with .zip or .whl
         if not (file.name.endswith(".zip") or file.name.endswith(".whl")):
@@ -169,16 +200,17 @@ def convert_wheels_to_tar_gz(dist = Path("dist")):
                 # create tar.gz file from tempdir
                 with tarfile.open(new_file, "w:gz") as tar:
                     tar.add(tempdir, arcname="./")
-        os.remove(file)
+        file.unlink()
         package["file_name"] = new_file.name
         # update sha256 hash
         new_file_bytes = new_file.read_bytes()
         new_file_hash = hashlib.sha256(new_file_bytes).hexdigest()
         package["sha256"] = new_file_hash
-    
+
     with open(dist / "pyodide-lock.json", "w") as file:
         json.dump(lock, file)
 
+
 if __name__ == "__main__":
     if len(sys.argv) != 2:
         print("Usage: python script.py <tag>")
@@ -187,11 +219,14 @@ def convert_wheels_to_tar_gz(dist = Path("dist")):
 
     with open("required_packages.txt", "r") as file:
         required_packages = file.read().split("\n")
-    status = os.system(f"pyodide build-recipes --install {' '.join(required_packages)}")
-    if status != 0:
-        raise Exception("Failed to build recipes")
-    
+    result = subprocess.run(
+        ["pyodide", "build-recipes", "--install", *required_packages]
+    )
+    if result.returncode != 0:
+        print("Failed to build recipes", file=sys.stderr)
+        sys.exit(result.returncode)
+
     convert_wheels_to_tar_gz()
-    
+
     make_bundle(tag)
-    upload_to_r2(tag)
+    upload_to_r2(tag)