diff --git a/src/nonebot_plugin_parser/download/__init__.py b/src/nonebot_plugin_parser/download/__init__.py index c60c1568..26734b46 100644 --- a/src/nonebot_plugin_parser/download/__init__.py +++ b/src/nonebot_plugin_parser/download/__init__.py @@ -1,20 +1,13 @@ import asyncio from pathlib import Path -from functools import partial -from contextlib import contextmanager from urllib.parse import urljoin import httpx import aiofiles import curl_cffi from nonebot import logger, get_driver -from rich.progress import ( - Progress, - BarColumn, - TextColumn, - DownloadColumn, -) +from .rich import progress_bar, add_progress_task from .task import auto_task from ..utils import merge_av, safe_unlink, generate_file_name, is_module_available from ..config import pconfig @@ -31,22 +24,6 @@ def __init__(self): async def aclose(self): await self.client.aclose() - @staticmethod - @contextmanager - def rich_progress( - desc: str, - total: int | None = None, - ): - with Progress( - TextColumn("[bold blue]{task.description}", justify="right"), - BarColumn(bar_width=None), - "[progress.percentage]{task.percentage:>3.1f}%", - "•", - DownloadColumn(), - ) as progress: - task_id = progress.add_task(description=desc, total=total) - yield partial(progress.update, task_id) - @staticmethod def _validate_content_length( response: httpx.Response | curl_cffi.Response, @@ -84,10 +61,11 @@ async def _download_file_with_httpx( response.raise_for_status() content_length = self._validate_content_length(response) - with self.rich_progress( - f"httpx | {file_path.name}", - content_length, - ) as update_progress: + with progress_bar: + update_progress = add_progress_task( + f"httpx | {file_path.name}", + content_length, + ) async with aiofiles.open(file_path, "wb") as file: async for chunk in response.aiter_bytes(chunk_size): await file.write(chunk) @@ -112,10 +90,11 @@ async def _download_file_with_curl_cffi( response.raise_for_status() content_length = self._validate_content_length(response) - with self.rich_progress( - f"curl_cffi | {file_path.name}", - content_length, - ) as update_progress: + with progress_bar: + update_progress = add_progress_task( + f"curl_cffi | {file_path.name}", + content_length, + ) async with aiofiles.open(file_path, "wb") as file: async for chunk in response.aiter_content(chunk_size=8192): await file.write(chunk) @@ -241,15 +220,14 @@ async def download_m3u8( video_path = pconfig.cache_dir / video_name try: - async with aiofiles.open(video_path, "wb") as f: - total_size = 0 - with self.rich_progress(desc=video_name) as update_progress: + with progress_bar: + async with aiofiles.open(video_path, "wb") as f: + update_progress = add_progress_task(desc=video_name) for url in await self._get_m3u8_slices(m3u8_url): async with self.client.stream("GET", url, headers=ext_headers) as response: async for chunk in response.aiter_bytes(chunk_size=1024 * 1024): await f.write(chunk) - total_size += len(chunk) - update_progress(advance=len(chunk), total=total_size) + update_progress(advance=len(chunk)) except httpx.HTTPError: await safe_unlink(video_path) logger.exception("m3u8 视频下载失败") diff --git a/src/nonebot_plugin_parser/download/rich.py b/src/nonebot_plugin_parser/download/rich.py new file mode 100644 index 00000000..7933b36e --- /dev/null +++ b/src/nonebot_plugin_parser/download/rich.py @@ -0,0 +1,25 @@ +from functools import partial + +from rich.progress import ( + Progress, + BarColumn, + TextColumn, + DownloadColumn, +) + +progress_bar: Progress = Progress( + TextColumn("[bold blue]{task.description}", justify="right"), + BarColumn(bar_width=None), + "[progress.percentage]{task.percentage:>3.1f}%", + "•", + DownloadColumn(), +) + + +def add_progress_task( + desc: str, + total: int | None = None, +): + task_id = progress_bar.add_task(description=desc, total=total) + progress_bar.start_task(task_id) + return partial(progress_bar.update, task_id) diff --git a/tests/parsers/test_acfun.py b/tests/parsers/test_acfun.py index ed5ccb33..25861e5f 100644 --- a/tests/parsers/test_acfun.py +++ b/tests/parsers/test_acfun.py @@ -8,6 +8,7 @@ async def test_parse(): # url = "https://www.acfun.cn/v/ac46593564" url = "https://www.acfun.cn/v/ac11348130" + # url = "https://www.acfun.cn/v/ac48729680" parser = AcfunParser() async def parse_acfun_url(url: str) -> None: