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