diff --git a/src/nonebot_plugin_parser/download/__init__.py b/src/nonebot_plugin_parser/download/__init__.py index 26734b46..0c1c4344 100644 --- a/src/nonebot_plugin_parser/download/__init__.py +++ b/src/nonebot_plugin_parser/download/__init__.py @@ -7,7 +7,7 @@ import curl_cffi from nonebot import logger, get_driver -from .rich import progress_bar, add_progress_task +from .rich import 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 @@ -61,15 +61,14 @@ async def _download_file_with_httpx( response.raise_for_status() content_length = self._validate_content_length(response) - 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) - update_progress(advance=len(chunk)) + 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) + update_progress(advance=len(chunk)) return file_path @@ -90,15 +89,14 @@ async def _download_file_with_curl_cffi( response.raise_for_status() content_length = self._validate_content_length(response) - 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) - update_progress(advance=len(chunk)) + 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) + update_progress(advance=len(chunk)) return file_path @@ -220,14 +218,13 @@ async def download_m3u8( video_path = pconfig.cache_dir / video_name try: - 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) - update_progress(advance=len(chunk)) + 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) + 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 index 7933b36e..0903e2b5 100644 --- a/src/nonebot_plugin_parser/download/rich.py +++ b/src/nonebot_plugin_parser/download/rich.py @@ -23,3 +23,18 @@ def add_progress_task( task_id = progress_bar.add_task(description=desc, total=total) progress_bar.start_task(task_id) return partial(progress_bar.update, task_id) + + +from nonebot import get_driver + +driver = get_driver() + + +@driver.on_startup +async def enter_progress_bar(): + progress_bar.start() + + +@driver.on_shutdown +async def exit_progress_bar(): + progress_bar.stop()