Port "Update pyodide py-compile command to accept list of files to exclude" (#9)

- pyodide/pyodide#4911

---------

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

Author: ryanking13

pyodide_build/_py_compile.py

@@ -11,8 +11,7 @@
 from packaging.tags import Tag
 from packaging.utils import parse_wheel_filename
 
-from pyodide_build.common import _get_sha256_checksum
-
+from .common import _get_sha256_checksum
 from .logger import logger, set_log_level
 
 
@@ -236,11 +235,7 @@ def _get_py_compiled_archive_name(path: Path) -> str | None:
     >>> import re
     >>> re.sub("cp[0-9]*", "cpxxx", _get_py_compiled_archive_name(Path("snowballstemmer-2.2.0-py2.py3-none-any.whl")))
     'snowballstemmer-2.2.0-cpxxx-none-any.whl'
-    >>> _get_py_compiled_archive_name(Path("test-1.0.0.zip"))
     """
-    # TODO: fix py-compilation of the following packages
-    if path.name.startswith(("RobotRaconteur", "astropy-", "opencv_python-")):
-        return None
 
     if path.suffix == ".whl":
         try:
@@ -249,9 +244,6 @@ def _get_py_compiled_archive_name(path: Path) -> str | None:
         except Exception as e:
             print(e)
             return None
-    elif path.name == "test-1.0.0.zip":
-        # We don't want to py-compile the test package
-        return None
     elif path.suffix == ".zip":
         # If it's a zip file with .py files, keep the same name
         with zipfile.ZipFile(path, "r") as zip_ref:
@@ -281,6 +273,7 @@ def _py_compile_archive_dir(
     keep: bool = True,
     verbose: bool = True,
     compression_level: int = 6,
+    excludes: list[str] | None = None,
 ) -> dict[str, str]:
     """Py-compile all wheels or zip files in a directory.
 
@@ -315,15 +308,21 @@ def _py_compile_archive_dir(
     for file_path in itertools.chain(
         *[input_dir.glob(ext) for ext in ["*.zip", "*.whl"]]
     ):
-        if (output_name := _get_py_compiled_archive_name(file_path)) is not None:
-            _compile(
-                file_path,
-                file_path.parent / output_name,
-                keep=keep,
-                verbose=verbose,
-                compression_level=compression_level,
-            )
-            name_mapping[file_path.name] = output_name
+        if excludes and any(file_path.name.startswith(exclude) for exclude in excludes):
+            continue
+
+        output_name = _get_py_compiled_archive_name(file_path)
+        if output_name is None:
+            continue
+
+        _compile(
+            file_path,
+            file_path.parent / output_name,
+            keep=keep,
+            verbose=verbose,
+            compression_level=compression_level,
+        )
+        name_mapping[file_path.name] = output_name
 
     lockfile_path = input_dir / "pyodide-lock.json"
     if name_mapping and lockfile_path.exists():

pyodide_build/cli/py_compile.py

@@ -1,8 +1,9 @@
+import re
 from pathlib import Path
 
 import typer
 
-from pyodide_build._py_compile import _py_compile_archive, _py_compile_archive_dir
+from .._py_compile import _py_compile_archive, _py_compile_archive_dir
 
 
 def main(
@@ -14,6 +15,10 @@ def main(
     compression_level: int = typer.Option(
         6, help="Compression level to use for the created zip file"
     ),
+    exclude: str = typer.Option(
+        "",
+        help="List of files to exclude from compilation, works only for directories. Defaults to no files.",
+    ),
 ) -> None:
     """Compile .py files to .pyc in a wheel, a zip file, or a folder with wheels or zip files.
 
@@ -24,6 +29,11 @@ def main(
         typer.echo(f"Error: {path} does not exist")
         raise typer.Exit(1)
 
+    # Convert the comma / space separated strings to lists
+    excludes = [
+        item.strip() for item in re.split(r",|\s", exclude) if item.strip() != ""
+    ]
+
     if path.is_file():
         if path.suffix not in [".whl", ".zip"]:
             typer.echo(
@@ -36,7 +46,11 @@ def main(
         )
     elif path.is_dir():
         _py_compile_archive_dir(
-            path, verbose=not silent, keep=keep, compression_level=compression_level
+            path,
+            verbose=not silent,
+            keep=keep,
+            compression_level=compression_level,
+            excludes=excludes,
         )
     else:
         typer.echo(f"{path=} is not a file or a directory")

pyodide_build/tests/test_cli.py

@@ -333,7 +333,11 @@ def test_py_compile(tmp_path, target, compression_level):
         target_path = wheel_path
 
     py_compile.main(
-        path=target_path, silent=False, keep=False, compression_level=compression_level
+        path=target_path,
+        silent=False,
+        keep=False,
+        compression_level=compression_level,
+        exclude="",
     )
     with zipfile.ZipFile(tmp_path / "python.zip", "r") as fh:
         if compression_level > 0:

pyodide_build/tests/test_py_compile.py

@@ -248,3 +248,30 @@ def test_py_compile_archive_dir(tmp_path, with_lockfile):
         "file_name": "some-path.tar",
         "checksum": "123",
     }
+
+
+@pytest.mark.parametrize("with_lockfile", [True, False])
+def test_py_compile_archive_dir_excludes(tmp_path, with_lockfile):
+    wheel_data = {
+        "a.so": "abc",
+        "b.txt": "123",
+        "METADATA": "a",
+        "packageB/a.py": "1+1",
+    }
+
+    _create_tmp_wheel(
+        "packageB", base_dir=tmp_path, data=wheel_data, tag="py3-none-any"
+    )
+
+    excludes = ["packageB-"]
+    mapping = _py_compile_archive_dir(tmp_path, keep=True, excludes=excludes)
+
+    assert mapping == {}
+
+    mapping2 = _py_compile_archive_dir(tmp_path, keep=False)
+
+    ver = sys.version_info
+    cpver = f"cp{ver.major}{ver.minor}"
+    assert mapping2 == {
+        "packageB-0.1.0-py3-none-any.whl": f"packageb-0.1.0-{cpver}-none-any.whl",
+    }