]> git.madduck.net Git - etc/vim.git/blob - gallery/gallery.py

madduck's git repository

Every one of the projects in this repository is available at the canonical URL git://git.madduck.net/madduck/pub/<projectpath> — see each project's metadata for the exact URL.

All patches and comments are welcome. Please squash your changes to logical commits before using git-format-patch and git-send-email to patches@git.madduck.net. If you'd read over the Git project's submission guidelines and adhered to them, I'd be especially grateful.

SSH access, as well as push access can be individually arranged.

If you use my repositories frequently, consider adding the following snippet to ~/.gitconfig and using the third clone URL listed for each project:

[url "git://git.madduck.net/madduck/"]
  insteadOf = madduck:

Refactor Jupyter magic handling (#2545)
[etc/vim.git] / 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()