Use pyodide-lock for pyodide-lock.json parsing in Python (#3949)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.pre-commit-config.yaml

@@ -85,6 +85,7 @@ repos:
           - resolvelib
           - rich
           - auditwheel_emscripten
+          - pyodide-lock==0.1.0a1
           - micropip
       - id: mypy
         name: mypy-tests

pyodide-build/pyodide_build/buildall.py

@@ -5,8 +5,6 @@ Build all of the packages in a given directory.
 """
 
 import dataclasses
-import hashlib
-import json
 import shutil
 import subprocess
 import sys
@@ -21,6 +19,8 @@ from threading import Lock, Thread
 from time import perf_counter, sleep
 from typing import Any
 
+from pyodide_lock import PyodideLockSpec
+from pyodide_lock.spec import PackageSpec as PackageLockSpec
 from rich.live import Live
 from rich.progress import BarColumn, Progress, TimeElapsedColumn
 from rich.spinner import Spinner
@@ -630,65 +630,54 @@ def build_from_graph(
                     build_queue.put((job_priority(dependent), dependent))
 
 
-def _generate_package_hash(full_path: Path) -> str:
-    sha256_hash = hashlib.sha256()
-    with open(full_path, "rb") as f:
-        while chunk := f.read(4096):
-            sha256_hash.update(chunk)
-    return sha256_hash.hexdigest()
-
-
 def generate_packagedata(
     output_dir: Path, pkg_map: dict[str, BasePackage]
-) -> dict[str, Any]:
-    packages: dict[str, Any] = {}
+) -> dict[str, PackageLockSpec]:
+    packages: dict[str, PackageLockSpec] = {}
     for name, pkg in pkg_map.items():
         if not pkg.file_name or pkg.package_type == "static_library":
             continue
         if not Path(output_dir, pkg.file_name).exists():
             continue
-        pkg_entry: Any = {
-            "name": name,
-            "version": pkg.version,
-            "file_name": pkg.file_name,
-            "install_dir": pkg.install_dir,
-            "sha256": _generate_package_hash(Path(output_dir, pkg.file_name)),
-            "package_type": pkg.package_type,
-            "imports": [],
-        }
+        pkg_entry = PackageLockSpec(
+            name=name,
+            version=pkg.version,
+            file_name=pkg.file_name,
+            install_dir=pkg.install_dir,
+            package_type=pkg.package_type,
+        )
+        pkg_entry.update_sha256(output_dir / pkg.file_name)
 
         pkg_type = pkg.package_type
         if pkg_type in ("shared_library", "cpython_module"):
             # We handle cpython modules as shared libraries
-            pkg_entry["shared_library"] = True
-            pkg_entry["install_dir"] = (
+            pkg_entry.shared_library = True
+            pkg_entry.install_dir = (
                 "stdlib" if pkg_type == "cpython_module" else "dynlib"
             )
 
-        pkg_entry["depends"] = [x.lower() for x in pkg.run_dependencies]
+        pkg_entry.depends = [x.lower() for x in pkg.run_dependencies]
 
         if pkg.package_type not in ("static_library", "shared_library"):
-            pkg_entry["imports"] = (
+            pkg_entry.imports = (
                 pkg.meta.package.top_level if pkg.meta.package.top_level else [name]
             )
 
         packages[name.lower()] = pkg_entry
 
         if pkg.unvendored_tests:
-            packages[name.lower()]["unvendored_tests"] = True
+            packages[name.lower()].unvendored_tests = True
 
             # Create the test package if necessary
-            pkg_entry = {
-                "name": name + "-tests",
-                "version": pkg.version,
-                "depends": [name.lower()],
-                "imports": [],
-                "file_name": pkg.unvendored_tests.name,
-                "install_dir": pkg.install_dir,
-                "sha256": _generate_package_hash(
-                    Path(output_dir, pkg.unvendored_tests.name)
-                ),
-            }
+            pkg_entry = PackageLockSpec(
+                name=name + "-tests",
+                version=pkg.version,
+                depends=[name.lower()],
+                file_name=pkg.unvendored_tests.name,
+                install_dir=pkg.install_dir,
+            )
+            pkg_entry.update_sha256(output_dir / pkg.unvendored_tests.name)
+
             packages[name.lower() + "-tests"] = pkg_entry
 
     # sort packages by name
@@ -698,7 +687,7 @@ def generate_packagedata(
 
 def generate_lockfile(
     output_dir: Path, pkg_map: dict[str, BasePackage]
-) -> dict[str, dict[str, Any]]:
+) -> PyodideLockSpec:
     """Generate the package.json file"""
 
     from . import __version__
@@ -713,7 +702,7 @@ def generate_lockfile(
         "python": sys.version.partition(" ")[0],
     }
     packages = generate_packagedata(output_dir, pkg_map)
-    return dict(info=info, packages=packages)
+    return PyodideLockSpec(info=info, packages=packages)
 
 
 def copy_packages_to_dist_dir(
@@ -810,9 +799,7 @@ def install_packages(
     logger.info(f"Writing pyodide-lock.json to {lockfile_path}")
 
     package_data = generate_lockfile(output_dir, pkg_map)
-    with lockfile_path.open("w") as fd:
-        json.dump(package_data, fd)
-        fd.write("\n")
+    package_data.to_json(lockfile_path)
 
 
 def set_default_build_args(build_args: BuildArgs) -> BuildArgs:

pyodide-build/pyodide_build/create_pypa_index.py

@@ -1,15 +1,11 @@
 from pathlib import Path
 from textwrap import dedent
-from typing import TypedDict
 
-
-class PackageInfo(TypedDict):
-    file_name: str
-    sha256: str
+from pyodide_lock.spec import PackageSpec
 
 
 def create_pypa_index(
-    packages: dict[str, PackageInfo], target_dir: Path, dist_url: str
+    packages: dict[str, PackageSpec], target_dir: Path, dist_url: str
 ) -> None:
     """Create a pip-compatible Python package (pypa) index to be used with a Pyodide virtual
     environment.
@@ -40,7 +36,7 @@ def create_pypa_index(
     packages = {
         pkgname: pkginfo
         for (pkgname, pkginfo) in packages.items()
-        if pkginfo["file_name"].endswith(".whl")
+        if pkginfo.file_name.endswith(".whl")
     }
     if not target_dir.exists():
         raise RuntimeError(f"target_dir={target_dir} does not exist")
@@ -87,8 +83,8 @@ def create_pypa_index(
 
     for pkgname, pkginfo in packages.items():
         pkgdir = index_dir / pkgname
-        filename = pkginfo["file_name"]
-        shasum = pkginfo["sha256"]
+        filename = pkginfo.file_name
+        shasum = pkginfo.sha256
         href = f"{dist_url}{filename}#sha256={shasum}"
         links_str = f'<a href="{href}">{pkgname}</a>\n'
         files_html = files_template.format(pkgname=pkgname, links=links_str)

pyodide-build/pyodide_build/install_xbuildenv.py

@@ -5,6 +5,8 @@ from pathlib import Path
 from urllib.error import HTTPError
 from urllib.request import urlopen, urlretrieve
 
+from pyodide_lock import PyodideLockSpec
+
 from . import build_env
 from .common import exit_with_stdio
 from .create_pypa_index import create_pypa_index
@@ -71,7 +73,7 @@ def install_xbuildenv(version: str, xbuildenv_path: Path) -> Path:
     cdn_base = f"https://cdn.jsdelivr.net/pyodide/v{version}/full/"
     lockfile_path = xbuildenv_root / "dist" / "pyodide-lock.json"
     if lockfile_path.exists():
-        lockfile_bytes = lockfile_path.read_bytes()
+        lockfile = PyodideLockSpec.from_json(lockfile_path)
     else:
         try:
             with urlopen(cdn_base + "pyodide-lock.json") as response:
@@ -80,9 +82,8 @@ def install_xbuildenv(version: str, xbuildenv_path: Path) -> Path:
             # Try again with old url
             with urlopen(cdn_base + "repodata.json") as response:
                 lockfile_bytes = response.read()
-    lockfile = json.loads(lockfile_bytes)
-    version = lockfile["info"]["version"]
-    create_pypa_index(lockfile["packages"], xbuildenv_root, cdn_base)
+        lockfile = PyodideLockSpec(**json.loads(lockfile_bytes))
+    create_pypa_index(lockfile.packages, xbuildenv_root, cdn_base)
 
     (xbuildenv_path / ".installed").touch()
 

pyodide-build/pyodide_build/tests/fixture.py

@@ -1,10 +1,9 @@
-import json
 import os
 import shutil
 from pathlib import Path
-from typing import Any
 
 import pytest
+from pyodide_lock import PyodideLockSpec
 
 from conftest import ROOT_PATH
 from pyodide_build import build_env
@@ -44,15 +43,16 @@ def temp_python_lib2(tmp_path_factory):
     yield libdir
 
 
-def mock_pyodide_lock() -> dict[str, Any]:
-    # TODO: use pydantic
-
-    return {
-        "info": {
+def mock_pyodide_lock() -> PyodideLockSpec:
+    return PyodideLockSpec(
+        info={
             "version": "0.22.1",
+            "arch": "wasm32",
+            "platform": "emscripten_xxx",
+            "python": "3.11",
         },
-        "packages": {},
-    }
+        packages={},
+    )
 
 
 @pytest.fixture(scope="module")
@@ -84,9 +84,7 @@ export HOSTSITEPACKAGES=$(PYODIDE_ROOT)/packages/.artifacts/lib/python$(PYMAJOR)
 """  # noqa: W191
     )
     (pyodide_root / "dist").mkdir()
-    (pyodide_root / "dist" / "pyodide-lock.json").write_text(
-        json.dumps(mock_pyodide_lock())
-    )
+    mock_pyodide_lock().to_json(pyodide_root / "dist" / "pyodide-lock.json")
 
     with chdir(base):
         archive_name = shutil.make_archive("xbuildenv", "tar")

pyodide-build/pyodide_build/tests/test_buildall.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from typing import Any
 
 import pytest
+from pyodide_lock.spec import PackageSpec
 
 from pyodide_build import buildall
 from pyodide_build.pywasmcross import BuildArgs
@@ -61,12 +62,10 @@ def test_generate_lockfile(tmp_path):
             hashes[pkg.name] = hashlib.sha256(f.read()).hexdigest()
 
     package_data = buildall.generate_lockfile(tmp_path, pkg_map)
-    assert set(package_data.keys()) == {"info", "packages"}
-    assert set(package_data["info"].keys()) == {"arch", "platform", "version", "python"}
-    assert package_data["info"]["arch"] == "wasm32"
-    assert package_data["info"]["platform"].startswith("emscripten")
+    assert package_data.info.arch == "wasm32"
+    assert package_data.info.platform.startswith("emscripten")
 
-    assert set(package_data["packages"]) == {
+    assert set(package_data.packages) == {
         "pkg_1",
         "pkg_1_1",
         "pkg_2",
@@ -74,22 +73,20 @@ def test_generate_lockfile(tmp_path):
         "pkg_3_1",
         "libtest_shared",
     }
-    assert package_data["packages"]["pkg_1"] == {
-        "name": "pkg_1",
-        "version": "1.0.0",
-        "file_name": "pkg_1.whl",
-        "depends": ["pkg_1_1", "pkg_3", "libtest_shared"],
-        "imports": ["pkg_1"],
-        "package_type": "package",
-        "install_dir": "site",
-        "sha256": hashes["pkg_1"],
-    }
-
-    assert (
-        package_data["packages"]["libtest_shared"]["package_type"] == "shared_library"
+    assert package_data.packages["pkg_1"] == PackageSpec(
+        name="pkg_1",
+        version="1.0.0",
+        file_name="pkg_1.whl",
+        depends=["pkg_1_1", "pkg_3", "libtest_shared"],
+        imports=["pkg_1"],
+        package_type="package",
+        install_dir="site",
+        sha256=hashes["pkg_1"],
     )
 
-    sharedlib_imports = package_data["packages"]["libtest_shared"]["imports"]
+    assert package_data.packages["libtest_shared"].package_type == "shared_library"
+
+    sharedlib_imports = package_data.packages["libtest_shared"].imports
     assert not sharedlib_imports, (
         "shared libraries should not have any imports, but got " f"{sharedlib_imports}"
     )

pyodide-build/pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "types-requests",
     "typer",
     "auditwheel-emscripten~=0.0.9",
+    "pyodide-lock==0.1.0a1",
     "resolvelib",
     "rich",
     "loky",

pyproject.toml

@@ -93,6 +93,7 @@ known-first-party = [
 ]
 known-third-party = [
   "build",
+  "pyodide_lock",
 ]
 
 [tool.ruff.mccabe]