convert.py

#!/usr/bin/env python3
import argparse
import os
import os.path
import shutil
import sys
import zipfile
import tarfile
import tempfile
import gzip
import bz2
import lzma
import zstd
import brotli


def current(wheel, output_dir):
    dirname = os.path.join(output_dir, "current")
    os.makedirs(dirname, exist_ok=True)
    shutil.copy(wheel, os.path.join(dirname, os.path.basename(wheel)))


def current_with_compression(name, compression):
    def format(wheel, output_dir):
        dirname = os.path.join(output_dir, name)
        os.makedirs(dirname, exist_ok=True)

        with zipfile.ZipFile(wheel, "r") as ofp:
            with zipfile.ZipFile(
                os.path.join(dirname, os.path.basename(wheel)), "w"
            ) as nfp:
                for info in ofp.infolist():
                    with ofp.open(info) as fp:
                        nfp.writestr(
                            info, fp.read(), compress_type=compression, compresslevel=9,
                        )

    return format


def tarball_with_compression(name, compressor):
    def format(wheel, output_dir):
        dirname = os.path.join(output_dir, name)
        os.makedirs(dirname, exist_ok=True)

        with zipfile.ZipFile(wheel, "r") as ofp:
            with zipfile.ZipFile(
                os.path.join(dirname, os.path.basename(wheel)), "w"
            ) as nfp:
                with tempfile.TemporaryFile() as tempfp:
                    # We want to collect everything that ISNT a .dist-info file,
                    # and wrap it into a tarfile. This could break if there are
                    # nested dist-info or something weird, but it's probably
                    # good enough for this test.
                    with tarfile.open(fileobj=tempfp, mode="w") as tfp:
                        for info in ofp.infolist():
                            if ".dist-info/" in info.filename:
                                continue

                            with ofp.open(info) as fp:
                                tinfo = tarfile.TarInfo(name=info.filename)
                                tinfo.size = info.file_size
                                tfp.addfile(tinfo, fp)

                    # Now we want to take that tarball, and compress it using our compressor
                    # TOOD: Compression
                    tempfp.seek(0)
                    nfp.writestr("data", compressor(tempfp.read()))

                # Now we want to just copy over the .dist-info files into the
                # new zip file, we're just going to hardcode with DEFLATE. We could
                # probably do something here with a different compression type, but
                # this is good enough for right now. Note that the RECORD file is
                # a good target for further optimizations.
                for info in ofp.infolist():
                    if ".dist-info/" in info.filename:
                        with ofp.open(info) as fp:
                            nfp.writestr(
                                info, fp.read(), compress_type=zipfile.ZIP_DEFLATED, compresslevel=9,
                            )

    return format


FORMATS = {
    "current": current,
    "current+deflate": current_with_compression("current+deflate", zipfile.ZIP_DEFLATED),
    "current+bzip2": current_with_compression("current+bzip2", zipfile.ZIP_BZIP2),
    "current+lzma": current_with_compression("current+lzma", zipfile.ZIP_LZMA),
    "tar+gz": tarball_with_compression("tar+gz", gzip.compress),
    "tar+bz2": tarball_with_compression("tar+bz2", bz2.compress),
    "tar+xz": tarball_with_compression("tar+xz", lzma.compress),
    "tar+zstd": tarball_with_compression("tar+zstd", lambda data: zstd.compress(data, 22)),
    "tar+brotli": tarball_with_compression("tar+brotli", lambda data: brotli.compress(data, quality=11)),
}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--format", "-f", action="append",
    )
    parser.add_argument("--output-dir", "-o", default="output")
    parser.add_argument("wheel")

    args = parser.parse_args(sys.argv[1:])

    formats = args.format
    if formats is None:
        formats = list(FORMATS)

    for format in formats:
        print(f"Converting to {format}")
        FORMATS[format](args.wheel, args.output_dir)


if __name__ == "__main__":
    main()