gallery/gallery.py

   1 import atexit
   2 import json
   3 import subprocess
   4 import tarfile
   5 import tempfile
   6 import traceback
   7 import venv
   8 import zipfile
   9 from argparse import ArgumentParser, Namespace
  10 from concurrent.futures import ThreadPoolExecutor
  11 from functools import lru_cache, partial
  12 from pathlib import Path
  13 from typing import Generator, List, NamedTuple, Optional, Tuple, Union, cast
  14 from urllib.request import urlopen, urlretrieve
  15
  16 PYPI_INSTANCE = "https://pypi.org/pypi"
  17 PYPI_TOP_PACKAGES = (
  18     "https://hugovk.github.io/top-pypi-packages/top-pypi-packages-30-days.min.json"
  19 )
  20 INTERNAL_BLACK_REPO = f"{tempfile.gettempdir()}/__black"
  21
  22 ArchiveKind = Union[tarfile.TarFile, zipfile.ZipFile]
  23
  24 subprocess.run = partial(subprocess.run, check=True)  # type: ignore
  25 # https://github.com/python/mypy/issues/1484
  26
  27
  28 class BlackVersion(NamedTuple):
  29     version: str
  30     config: Optional[str] = None
  31
  32
  33 def get_pypi_download_url(package: str, version: Optional[str]) -> str:
  34     with urlopen(PYPI_INSTANCE + f"/{package}/json") as page:
  35         metadata = json.load(page)
  36
  37     if version is None:
  38         sources = metadata["urls"]
  39     else:
  40         if version in metadata["releases"]:
  41             sources = metadata["releases"][version]
  42         else:
  43             raise ValueError(
  44                 f"No releases found with version ('{version}') tag. "
  45                 f"Found releases: {metadata['releases'].keys()}"
  46             )
  47
  48     for source in sources:
  49         if source["python_version"] == "source":
  50             break
  51     else:
  52         raise ValueError(f"Couldn't find any sources for {package}")
  53
  54     return cast(str, source["url"])
  55
  56
  57 def get_top_packages() -> List[str]:
  58     with urlopen(PYPI_TOP_PACKAGES) as page:
  59         result = json.load(page)
  60
  61     return [package["project"] for package in result["rows"]]
  62
  63
  64 def get_package_source(package: str, version: Optional[str]) -> str:
  65     if package == "cpython":
  66         if version is None:
  67             version = "main"
  68         return f"https://github.com/python/cpython/archive/{version}.zip"
  69     elif package == "pypy":
  70         if version is None:
  71             version = "branch/default"
  72         return (
  73             f"https://foss.heptapod.net/pypy/pypy/repository/{version}/archive.tar.bz2"
  74         )
  75     else:
  76         return get_pypi_download_url(package, version)
  77
  78
  79 def get_archive_manager(local_file: str) -> ArchiveKind:
  80     if tarfile.is_tarfile(local_file):
  81         return tarfile.open(local_file)
  82     elif zipfile.is_zipfile(local_file):
  83         return zipfile.ZipFile(local_file)
  84     else:
  85         raise ValueError("Unknown archive kind.")
  86
  87
  88 def get_first_archive_member(archive: ArchiveKind) -> str:
  89     if isinstance(archive, tarfile.TarFile):
  90         return archive.getnames()[0]
  91     elif isinstance(archive, zipfile.ZipFile):
  92         return archive.namelist()[0]
  93
  94
  95 def download_and_extract(package: str, version: Optional[str], directory: Path) -> Path:
  96     source = get_package_source(package, version)
  97
  98     local_file, _ = urlretrieve(source, directory / f"{package}-src")
  99     with get_archive_manager(local_file) as archive:
 100         archive.extractall(path=directory)
 101         result_dir = get_first_archive_member(archive)
 102     return directory / result_dir
 103
 104
 105 def get_package(
 106     package: str, version: Optional[str], directory: Path
 107 ) -> Optional[Path]:
 108     try:
 109         return download_and_extract(package, version, directory)
 110     except Exception:
 111         print(f"Caught an exception while downloading {package}.")
 112         traceback.print_exc()
 113         return None
 114
 115
 116 DEFAULT_SLICE = slice(None)  # for flake8
 117
 118
 119 def download_and_extract_top_packages(
 120     directory: Path,
 121     workers: int = 8,
 122     limit: slice = DEFAULT_SLICE,
 123 ) -> Generator[Path, None, None]:
 124     with ThreadPoolExecutor(max_workers=workers) as executor:
 125         bound_downloader = partial(get_package, version=None, directory=directory)
 126         for package in executor.map(bound_downloader, get_top_packages()[limit]):
 127             if package is not None:
 128                 yield package
 129
 130
 131 def git_create_repository(repo: Path) -> None:
 132     subprocess.run(["git", "init"], cwd=repo)
 133     git_add_and_commit(msg="Initial commit", repo=repo)
 134
 135
 136 def git_add_and_commit(msg: str, repo: Path) -> None:
 137     subprocess.run(["git", "add", "."], cwd=repo)
 138     subprocess.run(["git", "commit", "-m", msg, "--allow-empty"], cwd=repo)
 139
 140
 141 def git_switch_branch(
 142     branch: str, repo: Path, new: bool = False, from_branch: Optional[str] = None
 143 ) -> None:
 144     args = ["git", "checkout"]
 145     if new:
 146         args.append("-b")
 147     args.append(branch)
 148     if from_branch:
 149         args.append(from_branch)
 150     subprocess.run(args, cwd=repo)
 151
 152
 153 def init_repos(options: Namespace) -> Tuple[Path, ...]:
 154     options.output.mkdir(exist_ok=True)
 155
 156     if options.top_packages:
 157         source_directories = tuple(
 158             download_and_extract_top_packages(
 159                 directory=options.output,
 160                 workers=options.workers,
 161                 limit=slice(None, options.top_packages),
 162             )
 163         )
 164     else:
 165         source_directories = (
 166             download_and_extract(
 167                 package=options.pypi_package,
 168                 version=options.version,
 169                 directory=options.output,
 170             ),
 171         )
 172
 173     for source_directory in source_directories:
 174         git_create_repository(source_directory)
 175
 176     if options.black_repo is None:
 177         subprocess.run(
 178             ["git", "clone", "https://github.com/psf/black.git", INTERNAL_BLACK_REPO],
 179             cwd=options.output,
 180         )
 181         options.black_repo = options.output / INTERNAL_BLACK_REPO
 182
 183     return source_directories
 184
 185
 186 @lru_cache(8)
 187 def black_runner(version: str, black_repo: Path) -> Path:
 188     directory = tempfile.TemporaryDirectory()
 189     venv.create(directory.name, with_pip=True)
 190
 191     python = Path(directory.name) / "bin" / "python"
 192     subprocess.run([python, "-m", "pip", "install", "-e", black_repo])
 193
 194     atexit.register(directory.cleanup)
 195     return python
 196
 197
 198 def format_repo_with_version(
 199     repo: Path,
 200     from_branch: Optional[str],
 201     black_repo: Path,
 202     black_version: BlackVersion,
 203     input_directory: Path,
 204 ) -> str:
 205     current_branch = f"black-{black_version.version}"
 206     git_switch_branch(black_version.version, repo=black_repo)
 207     git_switch_branch(current_branch, repo=repo, new=True, from_branch=from_branch)
 208
 209     format_cmd: List[Union[Path, str]] = [
 210         black_runner(black_version.version, black_repo),
 211         (black_repo / "black.py").resolve(),
 212         ".",
 213     ]
 214     if black_version.config:
 215         format_cmd.extend(["--config", input_directory / black_version.config])
 216
 217     subprocess.run(format_cmd, cwd=repo, check=False)  # ensure the process
 218     # continuess to run even it can't format some files. Reporting those
 219     # should be enough
 220     git_add_and_commit(f"Format with black:{black_version.version}", repo=repo)
 221
 222     return current_branch
 223
 224
 225 def format_repos(repos: Tuple[Path, ...], options: Namespace) -> None:
 226     black_versions = tuple(
 227         BlackVersion(*version.split(":")) for version in options.versions
 228     )
 229
 230     for repo in repos:
 231         from_branch = None
 232         for black_version in black_versions:
 233             from_branch = format_repo_with_version(
 234                 repo=repo,
 235                 from_branch=from_branch,
 236                 black_repo=options.black_repo,
 237                 black_version=black_version,
 238                 input_directory=options.input,
 239             )
 240         git_switch_branch("main", repo=repo)
 241
 242     git_switch_branch("main", repo=options.black_repo)
 243
 244
 245 def main() -> None:
 246     parser = ArgumentParser(description="""Black Gallery is a script that
 247     automates the process of applying different Black versions to a selected
 248     PyPI package and seeing the results between versions.""")
 249
 250     group = parser.add_mutually_exclusive_group(required=True)
 251     group.add_argument("-p", "--pypi-package", help="PyPI package to download.")
 252     group.add_argument(
 253         "-t", "--top-packages", help="Top n PyPI packages to download.", type=int
 254     )
 255
 256     parser.add_argument("-b", "--black-repo", help="Black's Git repository.", type=Path)
 257     parser.add_argument(
 258         "-v",
 259         "--version",
 260         help=(
 261             "Version for given PyPI package. Will be discarded if used with -t option."
 262         ),
 263     )
 264     parser.add_argument(
 265         "-w",
 266         "--workers",
 267         help=(
 268             "Maximum number of threads to download with at the same time. "
 269             "Will be discarded if used with -p option."
 270         ),
 271     )
 272     parser.add_argument(
 273         "-i",
 274         "--input",
 275         default=Path("/input"),
 276         type=Path,
 277         help="Input directory to read configuration.",
 278     )
 279     parser.add_argument(
 280         "-o",
 281         "--output",
 282         default=Path("/output"),
 283         type=Path,
 284         help="Output directory to download and put result artifacts.",
 285     )
 286     parser.add_argument("versions", nargs="*", default=("main",), help="")
 287
 288     options = parser.parse_args()
 289     repos = init_repos(options)
 290     format_repos(repos, options)
 291
 292
 293 if __name__ == "__main__":
 294     main()