Honestly, I don't even *need* AsyncDownloader

This commit is contained in:
Casey 2024-05-22 14:37:27 +03:00
parent 66fa2a87f4
commit 21977244fd
Signed by: hkc
GPG Key ID: F0F6CFE11CDB0960
1 changed files with 2 additions and 57 deletions

59
make.py
View File

@ -1,11 +1,10 @@
#!/usr/bin/env python3
import json
from pathlib import Path
from typing import Optional, Union
from typing import Union
import asyncio
from httpx import AsyncClient
from asyncio.queues import Queue
from rich.progress import TaskID, Progress
from rich.progress import Progress
from zipfile import ZipFile
from PIL import Image
@ -43,60 +42,6 @@ meta_url: str = (
meta_filename = INPUT_PATH / f"meta-{meta_url.split('/')[-1]}"
class AsyncDownloader:
def __init__(
self, files: Optional[list[tuple[str, Path]]] = None, maxfiles: int = 8
):
self._queue: Queue[tuple[str, Path]] = Queue(maxfiles)
self._client = AsyncClient()
self._progress: Optional[Progress] = None
self._tasks: list[TaskID] = []
for url, path in files or []:
self._queue.put_nowait((url, path))
def add_file(self, url: str, path: Union[Path, str]):
self._queue.put_nowait((url, Path(path)))
async def __aenter__(self):
await self._client.__aenter__()
self._progress = Progress()
self._progress.__enter__()
for i in range(self._queue.maxsize):
self._tasks.append(
self._progress.add_task(f"Downloader {i + 1}", start=False)
)
return self
async def _worker(self, task: TaskID):
if not self._progress:
return
while not self._queue.empty():
url, path = await self._queue.get()
with path.open("wb") as fout:
async with self._client.stream("GET", url) as req:
size = int(req.headers.get("Content-Length", 0))
self._progress.start_task(task)
self._progress.update(
task,
total=size,
completed=0,
description=f"{path.name}",
)
async for chunk in req.aiter_bytes(8192):
self._progress.advance(task, fout.write(chunk))
self._progress.stop_task(task)
async def run(self):
await asyncio.gather(*[self._worker(task) for task in self._tasks])
async def __aexit__(self, a, b, c):
if self._progress is None:
raise ValueError("how did that happen?")
self._progress.__exit__(a, b, c)
await self._client.__aexit__()
async def download_if_missing(url: str, path: Union[Path, str]):
path = Path(path)
if path.exists():