bazel/pyc_wheel.py

# Copyright (c) 2016 Grant Patten
# Copyright (c) 2019-2021 Adam Karpierz
# Licensed under the MIT License
# https://opensource.org/licenses/MIT

"""Compile all py files in a wheel to pyc files."""

import sys
import os
import distutils
import re
import stat
import shutil
import tempfile
import glob
import compileall
import zipfile
import hashlib
import csv
import base64
from datetime import datetime
from pathlib import Path

__all__ = ('convert_wheel', 'main')


HASH_ALGORITHM = hashlib.sha256


def convert_wheel(whl_file: Path, *, exclude=None, with_backup=False, quiet=False):
    """Generate a new whl with only pyc files."""

    if whl_file.suffix != ".whl":
        raise TypeError("File to convert must be a *.whl")

    if exclude: exclude = re.compile(exclude)

    dist_info = "-".join(whl_file.stem.split("-")[:-3])

    whl_dir  = tempfile.mkdtemp()
    whl_path = Path(whl_dir)

    try:
        # Extract our zip file temporarily
        with zipfile.ZipFile(str(whl_file), "r") as whl_zip:
            whl_zip.extractall(whl_dir)
            members = [member for member in whl_zip.infolist()
                       if member.is_dir() or not member.filename.endswith(".py")]

        # Compile all py files
        if not compileall.compile_dir(whl_dir, rx=exclude,
                                      ddir="<{}>".format(dist_info),
                                      quiet=int(quiet), force=True, legacy=True):
            raise RuntimeError("Error compiling Python sources in wheel "
                               "{!s}".format(whl_file.name))

        # Remove all original py files
        for py_file in whl_path.glob("**/*.py"):
            if py_file.is_file():
                if exclude is None or not exclude.search(str(py_file)):
                    if not quiet: print("Deleting py file: {!s}".format(py_file))
                    py_file.chmod(stat.S_IWUSR)
                    py_file.unlink()

        for root, dirs, files in os.walk(whl_dir):
            for fname in files:
                if fname.endswith(".py"):
                    py_file = Path(root)/fname
                    if exclude is None or not exclude.search(str(py_file)):
                        if not quiet: print("Removing file: {!s}".format(py_file))
                        py_file.chmod(stat.S_IWUSR)
                        py_file.unlink()

        for member in members:
            file_path = whl_path/member.filename
            timestamp = datetime(*member.date_time).timestamp()
            try:
                os.utime(str(file_path), (timestamp, timestamp))
            except Exception:
                pass  # ignore errors

        # dist_info_path = whl_path/"{}.dist-info".format(dist_info)
        # rewrite_dist_info(dist_info_path, exclude=exclude)

        # Rezip the file with the new version info
        whl_file_zip = whl_path.with_suffix(".zip")
        if whl_file_zip.exists(): whl_file_zip.unlink()
        shutil.make_archive(whl_dir, "zip", root_dir=whl_dir)
        if with_backup:
            whl_file.replace(whl_file.with_suffix(whl_file.suffix + ".bak"))
        shutil.move(str(whl_file_zip), str(whl_file))
    finally:
        # Clean up original directory
        shutil.rmtree(whl_dir, ignore_errors=True)


def rewrite_dist_info(dist_info_path: Path, *, exclude=None):
    """Rewrite the record file with pyc files instead of py files."""

    whl_path = dist_info_path.resolve().parent

    # Rewrite the record file with pyc files instead of py files.

    record_path = dist_info_path/"RECORD"
    record_path.chmod(stat.S_IWUSR | stat.S_IRUSR)

    record_data = []
    with record_path.open("r") as record:
        for file_dest, file_hash, file_len in csv.reader(record):
            if file_dest.endswith(".py"):
                # Do not keep py files, replace with pyc files
                if exclude is None or not exclude.search(file_dest):
                    file_dest = Path(file_dest)
                    # import platform
                    # pyc_fname = "{}.{}-{}{}.pyc".format(
                    #             file_dest.stem,
                    #             platform.python_implementation().lower(),
                    #             sys.version_info.major,
                    #             sys.version_info.minor)
                    # pyc_file = file_dest.parent/"__pycache__"/pyc_fname
                    pyc_file = file_dest.with_suffix(".pyc")
                    file_dest = str(pyc_file)

                    pyc_path = whl_path/pyc_file
                    with pyc_path.open("rb") as f:
                        data = f.read()
                    file_hash = HASH_ALGORITHM(data)
                    file_hash = "{}={}".format(file_hash.name,
                                               _b64encode(file_hash.digest()))
                    file_len  = len(data)
            record_data.append((file_dest, file_hash, file_len))

    with record_path.open("w", newline="\n") as record:
        csv.writer(record,
                   lineterminator="\n").writerows(sorted(set(record_data)))

    # Rewrite the wheel info file.

    wheel_path = dist_info_path/"WHEEL"
    wheel_path.chmod(stat.S_IWUSR | stat.S_IRUSR)

    with wheel_path.open("r") as wheel:
        wheel_data = wheel.readlines()

    tags = [line.split(" ")[1].strip()
            for line in wheel_data if line.startswith("Tag: ")]
    if not tags:
        raise RuntimeError("No tags present in {}/{}; cannot determine target"
                           " wheel filename".format(wheel_path.parent.name,
                                                    wheel_path.name))
    # Reassemble the tags for the wheel file
    interps   = sorted({tag.split("-")[0] for tag in tags})
    abis      = sorted({tag.split("-")[1] for tag in tags})
    platforms = sorted({tag.split("-")[2] for tag in tags})

    # [...]

    with wheel_path.open("w") as wheel:
        wheel.writelines(wheel_data)


def _get_platform():
    """Return our platform name 'win32', 'linux_x86_64'"""
    result = distutils.util.get_platform().replace(".", "_").replace("-", "_")
    if result == "linux_x86_64" and sys.maxsize == 2147483647:
        # pip pull request #3497
        result = "linux_i686"
    return result


def _b64encode(data):
    """urlsafe_b64encode without padding"""
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("utf-8")


def main(argv=sys.argv[1:]):
    from argparse import ArgumentParser
    parser = ArgumentParser(description="Compile all py files in a wheel")
    parser.add_argument("whl_file",
                        help="Path (can contain wildcards) to whl(s) to convert")
    parser.add_argument("--exclude", default=None,
                        help="skip files matching the regular expression; "
                             "the regexp is searched for in the full path "
                             "of each file considered for compilation")
    parser.add_argument("--with_backup", default=False, action="store_true",
                        help="Indicates whether the backup will be created.")
    parser.add_argument("--quiet", default=False, action="store_true",
                        help="Indicates whether the filenames and other "
                             "conversion information will be printed to "
                             "the standard output.")
    args = parser.parse_args(argv)
    for whl_file in glob.iglob(args.whl_file):
        convert_wheel(Path(whl_file), exclude=args.exclude,
                      with_backup=args.with_backup, quiet=args.quiet)

if __name__ == '__main__':
    main()