download.py

import asyncio
import hashlib
import json
import time
from pathlib import Path
from typing import List, Tuple

import httpx
from rich.progress import Progress

from .config import meme_config
from .log import logger
from .version import __version__


def _resource_url(base_url: str, name: str) -> str:
    return f"{base_url}v{__version__}/{name}"


# https://github.com/mnixry/nonebot-plugin-gocqhttp/blob/main/nonebot_plugin_gocqhttp/process/download.py
async def get_fastest_mirror() -> List[str]:
    assert meme_config.resource.resource_urls, "No resource url specified."

    async def head_mirror(client: httpx.AsyncClient, base_url: str):
        begin_time = time.time()
        response = await client.head(
            _resource_url(base_url, "resources/fonts/NotoSansSC-Regular.otf"), timeout=5
        )
        response.raise_for_status()
        elapsed_time = (time.time() - begin_time) * 1000
        return {"base_url": base_url, "elapsed_time": elapsed_time}

    async with httpx.AsyncClient() as client:
        results = await asyncio.gather(
            *(
                head_mirror(client, domain)
                for domain in meme_config.resource.resource_urls
            ),
            return_exceptions=True,
        )
    results = sorted(
        (result for result in results if not isinstance(result, Exception)),
        key=lambda r: r["elapsed_time"],
    )
    return [result["base_url"] for result in results]


async def check_resources():
    semaphore = asyncio.Semaphore(10)

    available_urls = (
        [meme_config.resource.resource_url]
        if meme_config.resource.resource_url
        else (await get_fastest_mirror())
    )
    logger.debug(f"Available resource urls: {available_urls}")
    if not available_urls:
        logger.warning("No resource url available.")
        return

    async def _download(client: httpx.AsyncClient, name: str):
        async with semaphore:
            for base_url in available_urls:
                url = _resource_url(base_url, name)
                try:
                    resp = await client.get(url, timeout=20, follow_redirects=True)
                    resp.raise_for_status()
                    return resp.content
                except httpx.HTTPError:
                    pass
            logger.warning(f"{name} download failed！")

    async with httpx.AsyncClient() as client:
        if content := await _download(client, "resources/resource_list.json"):
            resource_list = json.loads(content.decode("utf-8"))
        else:
            return

    download_list: List[Tuple[Path, str]] = []
    for resource in resource_list:
        file_name = str(resource["path"])
        file_hash = str(resource["hash"])
        file_path = Path(__file__).parent / "memes" / file_name
        if (
            file_path.exists()
            and hashlib.md5(file_path.read_bytes()).hexdigest() == file_hash
        ):
            continue
        else:
            download_list.append((file_path, f"meme_generator/memes/{file_name}"))

    if len(download_list):
        logger.info("Downloading images ...")
    else:
        return

    async with httpx.AsyncClient() as client:

        async def download_image(file_path: Path, file_name: str):
            if content := await _download(client, file_name):
                file_path.parent.mkdir(parents=True, exist_ok=True)
                with file_path.open("wb") as f:
                    f.write(content)

        with Progress(
            *Progress.get_default_columns(), "[yellow]{task.completed}/{task.total}"
        ) as progress:
            progress_task = progress.add_task(
                "[green]Downloading...", total=len(download_list)
            )
            tasks = [
                download_image(file_path, file_name)
                for file_path, file_name in download_list
            ]
            for task in asyncio.as_completed(tasks):
                await task
                progress.update(progress_task, advance=1)