gallery/gallery.py

import atexit
import json
import subprocess
import tarfile
import tempfile
import traceback
import venv
import zipfile
from argparse import ArgumentParser, Namespace
from concurrent.futures import ThreadPoolExecutor
from functools import lru_cache, partial
from pathlib import Path
from typing import (  # type: ignore # typing can't see Literal
    Generator,
    List,
    Literal,
    NamedTuple,
    Optional,
    Tuple,
    Union,
    cast,
)
from urllib.request import urlopen, urlretrieve

PYPI_INSTANCE = "https://pypi.org/pypi"
PYPI_TOP_PACKAGES = (
    "https://hugovk.github.io/top-pypi-packages/top-pypi-packages-{days}-days.json"
)
INTERNAL_BLACK_REPO = f"{tempfile.gettempdir()}/__black"

ArchiveKind = Union[tarfile.TarFile, zipfile.ZipFile]
Days = Union[Literal[30], Literal[365]]

subprocess.run = partial(subprocess.run, check=True)  # type: ignore
# https://github.com/python/mypy/issues/1484


class BlackVersion(NamedTuple):
    version: str
    config: Optional[str] = None


def get_pypi_download_url(package: str, version: Optional[str]) -> str:
    with urlopen(PYPI_INSTANCE + f"/{package}/json") as page:
        metadata = json.load(page)

    if version is None:
        sources = metadata["urls"]
    else:
        if version in metadata["releases"]:
            sources = metadata["releases"][version]
        else:
            raise ValueError(
                f"No releases found with version ('{version}') tag. "
                f"Found releases: {metadata['releases'].keys()}"
            )

    for source in sources:
        if source["python_version"] == "source":
            break
    else:
        raise ValueError(f"Couldn't find any sources for {package}")

    return cast(str, source["url"])


def get_top_packages(days: Days) -> List[str]:
    with urlopen(PYPI_TOP_PACKAGES.format(days=days)) as page:
        result = json.load(page)

    return [package["project"] for package in result["rows"]]


def get_package_source(package: str, version: Optional[str]) -> str:
    if package == "cpython":
        if version is None:
            version = "main"
        return f"https://github.com/python/cpython/archive/{version}.zip"
    elif package == "pypy":
        if version is None:
            version = "branch/default"
        return (
            f"https://foss.heptapod.net/pypy/pypy/repository/{version}/archive.tar.bz2"
        )
    else:
        return get_pypi_download_url(package, version)


def get_archive_manager(local_file: str) -> ArchiveKind:
    if tarfile.is_tarfile(local_file):
        return tarfile.open(local_file)
    elif zipfile.is_zipfile(local_file):
        return zipfile.ZipFile(local_file)
    else:
        raise ValueError("Unknown archive kind.")


def get_first_archive_member(archive: ArchiveKind) -> str:
    if isinstance(archive, tarfile.TarFile):
        return archive.getnames()[0]
    elif isinstance(archive, zipfile.ZipFile):
        return archive.namelist()[0]


def download_and_extract(package: str, version: Optional[str], directory: Path) -> Path:
    source = get_package_source(package, version)

    local_file, _ = urlretrieve(source, directory / f"{package}-src")
    with get_archive_manager(local_file) as archive:
        archive.extractall(path=directory)
        result_dir = get_first_archive_member(archive)
    return directory / result_dir


def get_package(
    package: str, version: Optional[str], directory: Path
) -> Optional[Path]:
    try:
        return download_and_extract(package, version, directory)
    except Exception:
        print(f"Caught an exception while downloading {package}.")
        traceback.print_exc()
        return None


DEFAULT_SLICE = slice(None)  # for flake8


def download_and_extract_top_packages(
    directory: Path,
    days: Days = 365,
    workers: int = 8,
    limit: slice = DEFAULT_SLICE,
) -> Generator[Path, None, None]:
    with ThreadPoolExecutor(max_workers=workers) as executor:
        bound_downloader = partial(get_package, version=None, directory=directory)
        for package in executor.map(bound_downloader, get_top_packages(days)[limit]):
            if package is not None:
                yield package


def git_create_repository(repo: Path) -> None:
    subprocess.run(["git", "init"], cwd=repo)
    git_add_and_commit(msg="Initial commit", repo=repo)


def git_add_and_commit(msg: str, repo: Path) -> None:
    subprocess.run(["git", "add", "."], cwd=repo)
    subprocess.run(["git", "commit", "-m", msg, "--allow-empty"], cwd=repo)


def git_switch_branch(
    branch: str, repo: Path, new: bool = False, from_branch: Optional[str] = None
) -> None:
    args = ["git", "checkout"]
    if new:
        args.append("-b")
    args.append(branch)
    if from_branch:
        args.append(from_branch)
    subprocess.run(args, cwd=repo)


def init_repos(options: Namespace) -> Tuple[Path, ...]:
    options.output.mkdir(exist_ok=True)

    if options.top_packages:
        source_directories = tuple(
            download_and_extract_top_packages(
                directory=options.output,
                workers=options.workers,
                limit=slice(None, options.top_packages),
            )
        )
    else:
        source_directories = (
            download_and_extract(
                package=options.pypi_package,
                version=options.version,
                directory=options.output,
            ),
        )

    for source_directory in source_directories:
        git_create_repository(source_directory)

    if options.black_repo is None:
        subprocess.run(
            ["git", "clone", "https://github.com/psf/black.git", INTERNAL_BLACK_REPO],
            cwd=options.output,
        )
        options.black_repo = options.output / INTERNAL_BLACK_REPO

    return source_directories


@lru_cache(8)
def black_runner(version: str, black_repo: Path) -> Path:
    directory = tempfile.TemporaryDirectory()
    venv.create(directory.name, with_pip=True)

    python = Path(directory.name) / "bin" / "python"
    subprocess.run([python, "-m", "pip", "install", "-e", black_repo])

    atexit.register(directory.cleanup)
    return python


def format_repo_with_version(
    repo: Path,
    from_branch: Optional[str],
    black_repo: Path,
    black_version: BlackVersion,
    input_directory: Path,
) -> str:
    current_branch = f"black-{black_version.version}"
    git_switch_branch(black_version.version, repo=black_repo)
    git_switch_branch(current_branch, repo=repo, new=True, from_branch=from_branch)

    format_cmd: List[Union[Path, str]] = [
        black_runner(black_version.version, black_repo),
        (black_repo / "black.py").resolve(),
        ".",
    ]
    if black_version.config:
        format_cmd.extend(["--config", input_directory / black_version.config])

    subprocess.run(format_cmd, cwd=repo, check=False)  # ensure the process
    # continuess to run even it can't format some files. Reporting those
    # should be enough
    git_add_and_commit(f"Format with black:{black_version.version}", repo=repo)

    return current_branch


def format_repos(repos: Tuple[Path, ...], options: Namespace) -> None:
    black_versions = tuple(
        BlackVersion(*version.split(":")) for version in options.versions
    )

    for repo in repos:
        from_branch = None
        for black_version in black_versions:
            from_branch = format_repo_with_version(
                repo=repo,
                from_branch=from_branch,
                black_repo=options.black_repo,
                black_version=black_version,
                input_directory=options.input,
            )
        git_switch_branch("main", repo=repo)

    git_switch_branch("main", repo=options.black_repo)


def main() -> None:
    parser = ArgumentParser(
        description="""Black Gallery is a script that
    automates the process of applying different Black versions to a selected
    PyPI package and seeing the results between versions."""
    )

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("-p", "--pypi-package", help="PyPI package to download.")
    group.add_argument(
        "-t", "--top-packages", help="Top n PyPI packages to download.", type=int
    )

    parser.add_argument("-b", "--black-repo", help="Black's Git repository.", type=Path)
    parser.add_argument(
        "-v",
        "--version",
        help=(
            "Version for given PyPI package. Will be discarded if used with -t option."
        ),
    )
    parser.add_argument(
        "-w",
        "--workers",
        help=(
            "Maximum number of threads to download with at the same time. "
            "Will be discarded if used with -p option."
        ),
    )
    parser.add_argument(
        "-i",
        "--input",
        default=Path("/input"),
        type=Path,
        help="Input directory to read configuration.",
    )
    parser.add_argument(
        "-o",
        "--output",
        default=Path("/output"),
        type=Path,
        help="Output directory to download and put result artifacts.",
    )
    parser.add_argument("versions", nargs="*", default=("main",), help="")

    options = parser.parse_args()
    repos = init_repos(options)
    format_repos(repos, options)


if __name__ == "__main__":
    main()