Add CLI command to py-compile wheels (#3253)

Co-authored-by: Gyeongjae Choi <def6488@gmail.com>

docs/project/changelog.md

@@ -15,6 +15,12 @@ myst:
 
 ## Unreleased
 
+### Build System
+
+- Added `pyodide py-compile` CLI command that py compiles a wheel, converting .py files
+  to .pyc files
+  {pr}`3253`
+
 ## Version 0.22.0
 
 _January 3, 2023_

packages/packaging/meta.yaml

@@ -4,8 +4,12 @@ package:
   top-level:
     - packaging
 source:
-  sha256: ef103e05f519cdc783ae24ea4e2e0f508a9c99b2d4969652eed6a2e1ea5bd522
-  url: https://files.pythonhosted.org/packages/05/8e/8de486cbd03baba4deef4142bd643a3e7bbe954a784dc1bb17142572d127/packaging-21.3-py3-none-any.whl
+  sha256: dd47c42927d89ab911e606518907cc2d3a1f38bbd026385970643f9c5b8ecfeb
+  url: https://files.pythonhosted.org/packages/df/9e/d1a7217f69310c1db8fdf8ab396229f55a699ce34a203691794c5d1cad0c/packaging-21.3.tar.gz
+
+  patches:
+    - patches/allow-cpp310-none-any-tag.patch
+
 requirements:
   run:
     - pyparsing

packages/packaging/patches/allow-cpp310-none-any-tag.patch

@@ -0,0 +1,54 @@
+From b8b2f8aa47c0676628f336487a31f602bea0c857 Mon Sep 17 00:00:00 2001
+From: Thimo <thimo.kraemer@joonis.de>
+Date: Sat, 20 Aug 2022 00:10:48 +0200
+Subject: [PATCH] Add a "cpNNN-none-any" tag (#541)
+
+Note: Can be removed with >=22.1 packaging release
+
+Closes #511
+
+See https://github.com/pypa/pip/issues/10923 for motivation.
+
+Co-authored-by: Brett Cannon <brett@python.org>
+---
+ packaging/tags.py  |  7 +++++--
+ tests/test_tags.py | 13 +++++++++++++
+ 2 files changed, 18 insertions(+), 2 deletions(-)
+
+diff --git a/packaging/tags.py b/packaging/tags.py
+index 5b6c5ffd..a0e1ea23 100644
+--- a/packaging/tags.py
++++ b/packaging/tags.py
+@@ -499,6 +499,9 @@ def sys_tags(*, warn: bool = False) -> Iterator[Tag]:
+         yield from generic_tags()
+
+     if interp_name == "pp":
+-        yield from compatible_tags(interpreter="pp3")
++        interp = "pp3"
++    elif interp_name == "cp":
++        interp = "cp" + interpreter_version(warn=warn)
+     else:
+-        yield from compatible_tags()
++        interp = None
++    yield from compatible_tags(interpreter=interp)
+diff --git a/tests/test_tags.py b/tests/test_tags.py
+index 06cd3b4a..39515e8d 100644
+--- a/tests/test_tags.py
++++ b/tests/test_tags.py
+@@ -1226,3 +1226,16 @@ def test_pypy_first_none_any_tag(self, monkeypatch):
+                 break
+
+         assert tag == tags.Tag("pp3", "none", "any")
++
++    def test_cpython_first_none_any_tag(self, monkeypatch):
++        # When building the complete list of cpython tags, make sure the first
++        # <interpreter>-none-any one is cpxx-none-any
++        monkeypatch.setattr(tags, "interpreter_name", lambda: "cp")
++
++        # Find the first tag that is ABI- and platform-agnostic.
++        for tag in tags.sys_tags():
++            if tag.abi == "none" and tag.platform == "any":
++                break
++
++        interpreter = f"cp{tags.interpreter_version()}"
++        assert tag == tags.Tag(interpreter, "none", "any")

pyodide-build/pyodide_build/_py_compile.py

@@ -0,0 +1,169 @@
+import py_compile
+import shutil
+import sys
+import zipfile
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+from packaging.tags import Tag
+from packaging.utils import parse_wheel_filename
+
+
+def _specialize_convert_tags(tags: set[Tag] | frozenset[Tag], wheel_name: str) -> Tag:
+    """Convert a sequence of wheel tags to a single tag corresponding
+    to the current interpreter and compatible with the py -> pyc compilation.
+
+    Having more than one output tag is not supported.
+
+    Examples
+    --------
+    >>> from packaging.tags import parse_tag
+    >>> tags = parse_tag("py2.py3-none-any")
+    >>> str(_specialize_convert_tags(set(tags), ""))
+    'cp310-none-any'
+    >>> tags = parse_tag("cp310-cp310-emscripten_3_1_24_wasm32")
+    >>> str(_specialize_convert_tags(set(tags), ""))
+    'cp310-cp310-emscripten_3_1_24_wasm32'
+    >>> tags = parse_tag("py310.py311-any-none")
+    >>> str(_specialize_convert_tags(set(tags), ""))
+    'cp310-any-none'
+    >>> tags = parse_tag("py36-abi3-none")
+    >>> str(_specialize_convert_tags(set(tags), ""))
+    'cp310-abi3-none'
+    """
+    if len(tags) == 0:
+        raise ValueError("Failed to parse tags from the wheel file name: {wheel_name}!")
+
+    output_tags = set()
+    interpreter = "cp" + "".join(str(el) for el in sys.version_info[:2])
+    for tag in tags:
+        output_tags.add(
+            Tag(interpreter=interpreter, abi=tag.abi, platform=tag.platform)
+        )
+
+    if len(output_tags) > 1:
+        # See https://github.com/pypa/packaging/issues/616
+        raise NotImplementedError(
+            "Found more than one output tag after py-compilation: "
+            f"{[str(tag) for tag in output_tags]} in {wheel_name}"
+        )
+
+    return list(output_tags)[0]
+
+
+def _py_compile_wheel_name(wheel_name: str) -> str:
+    """Return the name of the py-compiled wheel
+
+    See https://peps.python.org/pep-0427/ for more information.
+
+    Examples
+    --------
+    >>> _py_compile_wheel_name('micropip-0.1.0-py3-none-any.whl')
+    'micropip-0.1.0-cp310-none-any.whl'
+    >>> _py_compile_wheel_name("numpy-1.22.4-cp310-cp310-emscripten_3_1_24_wasm32.whl")
+    'numpy-1.22.4-cp310-cp310-emscripten_3_1_24_wasm32.whl'
+    >>> # names with '_' are preserved (instead of using '-')
+    >>> _py_compile_wheel_name("a_b-0.0.0-cp310-cp310-emscripten_3_1_24_wasm32.whl")
+    'a_b-0.0.0-cp310-cp310-emscripten_3_1_24_wasm32.whl'
+    >>> # if there are multiple tags (e.g. py2 & py3), we only keep the relevant one
+    >>> _py_compile_wheel_name('attrs-21.4.0-py2.py3-none-any.whl')
+    'attrs-21.4.0-cp310-none-any.whl'
+
+
+    # >>> msg = "Processing more than one tag is not implemented"
+    # >>> with pytest.rases(NotImplementedError, match=msg):
+    # ...     _py_compile_wheel_name("numpy-1.23.4-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl")
+    """
+    name, version, build, tags = parse_wheel_filename(wheel_name)
+    if build:
+        # TODO: not sure what to do here, but we never have such files in Pyodide
+        # Opened https://github.com/pypa/packaging/issues/616 about it.
+        raise NotImplementedError(f"build tag {build} not implemented")
+    output_name = f"{name.replace('-', '_')}-{version}-"
+    output_name += str(_specialize_convert_tags(tags, wheel_name=wheel_name))
+    return output_name + ".whl"
+
+
+def _py_compile_wheel(
+    wheel_path: Path,
+    keep: bool = True,
+    verbose: bool = True,
+) -> Path:
+    """Compile .py files to .pyc in a wheel
+
+    All non Python files are kept unchanged.
+
+    Parameters
+    ----------
+    wheel_path
+        input wheel path
+    keep
+        if False, delete the input file. Otherwise, it will be either kept or
+        renamed with a suffix .whl.old (if the input path == computed output
+        path)
+    verbose
+        print logging information
+
+    Returns
+    -------
+    wheel_path_out
+        processed wheel with .pyc files.
+
+
+    """
+    if wheel_path.suffix != ".whl":
+        raise ValueError(f"Error: only .whl files are supported, got {wheel_path.name}")
+
+    if not wheel_path.exists():
+        raise FileNotFoundError(f"{wheel_path} does not exist!")
+
+    wheel_name_out = _py_compile_wheel_name(wheel_path.name)
+    wheel_path_out = wheel_path.parent / wheel_name_out
+    if verbose:
+        print(f" - Running py-compile on {wheel_path} -> ", end="", flush=True)
+
+    with zipfile.ZipFile(wheel_path) as fh_zip_in, TemporaryDirectory() as temp_dir_str:
+        temp_dir = Path(temp_dir_str)
+        wheel_path_tmp = temp_dir / wheel_name_out
+        with zipfile.ZipFile(
+            wheel_path_tmp, mode="w", compression=zipfile.ZIP_DEFLATED
+        ) as fh_zip_out:
+            for name in fh_zip_in.namelist():
+                if name.endswith(".pyc"):
+                    # We are going to re-compile all .pyc files
+                    continue
+
+                stream = fh_zip_in.read(name)
+                if not name.endswith(".py"):
+                    # Write file without changes
+                    fh_zip_out.writestr(name, stream)
+                    continue
+
+                # Otherwise write file to disk and run py_compile
+                # Unfortunately py_compile doesn't support bytes input/output, it has to be real files
+                tmp_path_py = temp_dir / name.replace("/", "_")
+                tmp_path_py.write_bytes(stream)
+
+                tmp_path_pyc = temp_dir / (tmp_path_py.name + "c")
+                py_compile.compile(
+                    str(tmp_path_py), cfile=str(tmp_path_pyc), dfile=name, doraise=True
+                )
+
+                fh_zip_out.writestr(name + "c", tmp_path_pyc.read_bytes())
+        if wheel_name_out == wheel_path.name:
+            if keep:
+                if verbose:
+                    print(
+                        " (adding .old prefix to avoid overwriting input file) ->",
+                        end="",
+                        flush=True,
+                    )
+                wheel_path.rename(wheel_path.with_suffix(".whl.old"))
+        elif not keep:
+            # Remove input file
+            wheel_path.unlink()
+
+        shutil.copyfile(wheel_path_tmp, wheel_path_out)
+        if verbose:
+            print(f" {wheel_path_out.name}")
+    return wheel_path_out

pyodide-build/pyodide_build/cli/py_compile.py

@@ -0,0 +1,16 @@
+import sys
+from pathlib import Path
+
+import typer  # type: ignore[import]
+
+from pyodide_build._py_compile import _py_compile_wheel
+
+
+def main(
+    wheel_path: Path = typer.Argument(..., help="Path to the input wheel"),
+) -> None:
+    """Compile .py files to .pyc in a wheel"""
+    if wheel_path.suffix != ".whl":
+        typer.echo(f"Error: only .whl files are supported, got {wheel_path.name}")
+        sys.exit(1)
+    _py_compile_wheel(wheel_path, verbose=False)

pyodide-build/pyodide_build/common.py

@@ -104,6 +104,8 @@ def pyodide_tags() -> Iterator[Tag]:
     python_version = (int(PYMAJOR), int(PYMINOR))
     yield from cpython_tags(platforms=[PLATFORM], python_version=python_version)
     yield from compatible_tags(platforms=[PLATFORM], python_version=python_version)
+    # Following line can be removed once packaging 22.0 is released and we update to it.
+    yield Tag(interpreter=f"cp{PYMAJOR}{PYMINOR}", abi="none", platform="any")
 
 
 def find_matching_wheels(wheel_paths: Iterable[Path]) -> Iterator[Path]:

pyodide-build/pyodide_build/tests/test_common.py

@@ -110,6 +110,7 @@ def test_wheel_paths():
         f"py2.py3-none-{PLATFORM}",
         "py3-none-any",
         "py2.py3-none-any",
+        f"{current_version}-none-any",
     ]