gen_packed_apilevels.py

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import argparse
import base64
import io
import logging
import os
import shutil
import zipfile
from collections import namedtuple


try:
    import lzma  # noqa(F401)
    import tarfile

    has_tar_lzma = True
except ImportError:
    has_tar_lzma = False


Args = namedtuple("Args", ["inputs", "output", "tarxz"])


def parse_args():
    parser = argparse.ArgumentParser(
        description="Generate packed API levels python module"
    )

    parser.add_argument("api_files", nargs="+", help="API files")
    parser.add_argument(
        "-o",
        "--out",
        nargs=1,
        type=os.path.realpath,
        help="Generated python wrapper",
    )
    parser.add_argument(
        "--force-zip",
        action="store_true",
        help="Force the use of zip, even when tar and lzma are available",
    )

    args = parser.parse_args()

    global has_tar_lzma
    return Args(args.api_files, args.out[0], has_tar_lzma and not args.force_zip)


def compress_zip(inputs):
    logging.info("Compressing as zip")
    buf = io.BytesIO(b"")
    with zipfile.ZipFile(buf, "w") as zf:
        for input in inputs:
            logging.info("Adding %s", input)
            with open(input, "rb") as f:
                # Files are multi-MB, try to do it smartly.
                with zf.open(os.path.basename(input), mode="w") as f2:
                    shutil.copyfileobj(f, f2)
    buf.seek(0)
    return buf


def compress_tar_xz(inputs):
    logging.info("Compressing as tar.xz")
    buf = io.BytesIO(b"")
    tar = tarfile.open(fileobj=buf, mode="w:xz")

    for input in inputs:
        logging.info("Adding %s", input)
        # In case the inputs are symlinks, it's better to work with the file
        # object outright.
        with open(input, "rb") as f:
            info = tar.gettarinfo(arcname=os.path.basename(input), fileobj=f)
            tar.addfile(info, fileobj=f)

    tar.close()
    buf.seek(0)
    return buf


def compress_and_base_64(inputs, tar_xz):
    with compress_tar_xz(inputs) if tar_xz else compress_zip(inputs) as buf:
        return base64.b64encode(buf.getbuffer())


_FILE_TEMPLATE = """
#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import base64
import io
import re

{extra_imports}


_BASE64BLOB = "{base64_blob}"


{compression_specific_api}


def get_api_level_file(level):
    name = f"framework_classes_api_{{level}}.txt"
    return _load(name)


def get_api_levels():
    name_re = re.compile(r"^framework_classes_api_(\\d+)\\.txt$")
    return {{
        int(match.group(1)) for name in _all() for match in [name_re.match(name)] if match
    }}
"""


_TAR_XZ_IMPORTS = """
import lzma  # noqa(F401)
import tarfile
"""
_TAR_XZ_API = """
_TAR = tarfile.open(mode="r:xz", fileobj=io.BytesIO(base64.b64decode(_BASE64BLOB)))


def _load(name):
    global _TAR
    return _TAR.extractfile(name).read()


def _all():
    global _TAR
    return _TAR.getnames()
"""


_ZIP_IMPORTS = "import zipfile"
_ZIP_API = """
_ZIP = zipfile.ZipFile(io.BytesIO(base64.b64decode(_BASE64BLOB)), "r")


def _load(name):
    global _ZIP
    return _ZIP.read(name)


def _all():
    global _ZIP
    return _ZIP.namelist()
"""


def write_py_wrapper(base_64_bytes_blob, filename, tar_xz):
    base64_str = base_64_bytes_blob.decode("ascii")
    with open(filename, "w") as f:
        f.write(
            _FILE_TEMPLATE.format(
                extra_imports=_TAR_XZ_IMPORTS if tar_xz else _ZIP_IMPORTS,
                base64_blob=base64_str,
                compression_specific_api=_TAR_XZ_API if tar_xz else _ZIP_API,
            )
        )


def main():
    args = parse_args()
    base64_blob = compress_and_base_64(args.inputs, args.tarxz)
    write_py_wrapper(base64_blob, args.output, args.tarxz)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    main()