pyodide/pyodide-build/pyodide_build/_py_compile.py

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

from .logger import logger, set_log_level


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), ""))
    'cp311-none-any'
    >>> tags = parse_tag("cp311-cp311-emscripten_3_1_24_wasm32")
    >>> str(_specialize_convert_tags(set(tags), ""))
    'cp311-cp311-emscripten_3_1_24_wasm32'
    >>> tags = parse_tag("py310.py311-any-none")
    >>> str(_specialize_convert_tags(set(tags), ""))
    'cp311-any-none'
    >>> tags = parse_tag("py36-abi3-none")
    >>> str(_specialize_convert_tags(set(tags), ""))
    'cp311-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-cp311-none-any.whl'
    >>> _py_compile_wheel_name("numpy-1.22.4-cp311-cp311-emscripten_3_1_24_wasm32.whl")
    'numpy-1.22.4-cp311-cp311-emscripten_3_1_24_wasm32.whl'
    >>> # names with '_' are preserved (instead of using '-')
    >>> _py_compile_wheel_name("a_b-0.0.0-cp311-cp311-emscripten_3_1_24_wasm32.whl")
    'a_b-0.0.0-cp311-cp311-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-cp311-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-cp311-cp311-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 _compile(
    input_path: Path,
    output_path: Path,
    keep: bool = True,
    verbose: bool = True,
    compression_level: int = 6,
) -> None:
    """Compile all .py files in the zip archive to .pyc files.

    Parameters
    ----------
    input_path
        Path to the input archive.
    output_path
        Path to the output archive.
    compression_level
        Level of zip compression to apply. 0 means no compression. If a strictly
        positive integer is provided, ZIP_DEFLATED option is used.
    """
    output_name = output_path.name

    with set_log_level(logger, verbose):
        logger.debug(f"Running py-compile on {input_path} to {output_path}")

        if compression_level > 0:
            compression = zipfile.ZIP_DEFLATED
        else:
            compression = zipfile.ZIP_STORED

        with zipfile.ZipFile(
            input_path
        ) as fh_zip_in, TemporaryDirectory() as temp_dir_str:
            temp_dir = Path(temp_dir_str)
            output_path_tmp = temp_dir / output_name
            with zipfile.ZipFile(
                output_path_tmp,
                mode="w",
                compression=compression,
                compresslevel=compression_level,
            ) 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 output_path == input_path:
                if keep:
                    logger.debug("Adding .old suffix to avoid overwriting input file.")

                    backup_path = input_path.with_suffix(input_path.suffix + ".old")
                    input_path.rename(backup_path)
            elif not keep:
                # Remove input file
                input_path.unlink()

            shutil.copyfile(output_path_tmp, output_path)


def _py_compile_wheel(
    wheel_path: Path,
    keep: bool = True,
    verbose: bool = True,
    compression_level: int = 6,
) -> 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
    compression_level
        Level of zip compression to apply. 0 means no compression. If a strictly
        positive integer is provided, ZIP_DEFLATED option is used.

    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

    _compile(wheel_path, wheel_path_out, keep=keep, verbose=verbose)

    return wheel_path_out