From e60bb3787f3b7fa259ddfda61942388f5aa07ff3 Mon Sep 17 00:00:00 2001 From: fllesser Date: Sun, 2 Aug 2026 15:38:14 +0800 Subject: [PATCH 1/5] fix(download): fixed progress bar display with rich --- .../download/__init__.py | 45 ++++++++++--------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/src/nonebot_plugin_parser/download/__init__.py b/src/nonebot_plugin_parser/download/__init__.py index c60c1568..c1a84451 100644 --- a/src/nonebot_plugin_parser/download/__init__.py +++ b/src/nonebot_plugin_parser/download/__init__.py @@ -1,7 +1,6 @@ import asyncio from pathlib import Path from functools import partial -from contextlib import contextmanager from urllib.parse import urljoin import httpx @@ -27,25 +26,26 @@ def __init__(self): self.headers: dict[str, str] = COMMON_HEADER.copy() self.cache_dir: Path = pconfig.cache_dir self.client: httpx.AsyncClient = httpx.AsyncClient(timeout=DOWNLOAD_TIMEOUT, verify=False) + self.progress_bar: Progress = Progress( + TextColumn("[bold blue]{task.description}", justify="right"), + BarColumn(bar_width=None), + "[progress.percentage]{task.percentage:>3.1f}%", + "•", + DownloadColumn(), + ) async def aclose(self): await self.client.aclose() + self.progress_bar.stop() - @staticmethod - @contextmanager def rich_progress( + self, 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) + task_id = self.progress_bar.add_task(description=desc, total=total) + self.progress_bar.start_task(task_id) + return partial(self.progress_bar.update, task_id) @staticmethod def _validate_content_length( @@ -84,10 +84,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 self.progress_bar: + update_progress = self.rich_progress( + 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 +113,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 self.progress_bar: + update_progress = self.rich_progress( + 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) @@ -243,7 +245,8 @@ async def download_m3u8( try: async with aiofiles.open(video_path, "wb") as f: total_size = 0 - with self.rich_progress(desc=video_name) as update_progress: + with self.progress_bar: + update_progress = self.rich_progress(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): From 61d86cfd51b1c51b5ce8e68f93d5f2cfa35106ec Mon Sep 17 00:00:00 2001 From: fllesser Date: Sun, 2 Aug 2026 15:39:30 +0800 Subject: [PATCH 2/5] tweak --- src/nonebot_plugin_parser/download/__init__.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/nonebot_plugin_parser/download/__init__.py b/src/nonebot_plugin_parser/download/__init__.py index c1a84451..9c43ef4e 100644 --- a/src/nonebot_plugin_parser/download/__init__.py +++ b/src/nonebot_plugin_parser/download/__init__.py @@ -38,7 +38,7 @@ async def aclose(self): await self.client.aclose() self.progress_bar.stop() - def rich_progress( + def add_progress_task( self, desc: str, total: int | None = None, @@ -85,7 +85,7 @@ async def _download_file_with_httpx( content_length = self._validate_content_length(response) with self.progress_bar: - update_progress = self.rich_progress( + update_progress = self.add_progress_task( f"httpx | {file_path.name}", content_length, ) @@ -114,7 +114,7 @@ async def _download_file_with_curl_cffi( content_length = self._validate_content_length(response) with self.progress_bar: - update_progress = self.rich_progress( + update_progress = self.add_progress_task( f"curl_cffi | {file_path.name}", content_length, ) @@ -246,7 +246,7 @@ async def download_m3u8( async with aiofiles.open(video_path, "wb") as f: total_size = 0 with self.progress_bar: - update_progress = self.rich_progress(desc=video_name) + update_progress = self.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): From 3612d97083a594cf3e768ae24d22323adb1dbb4f Mon Sep 17 00:00:00 2001 From: fllesser Date: Sun, 2 Aug 2026 15:41:34 +0800 Subject: [PATCH 3/5] tweak --- src/nonebot_plugin_parser/download/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/nonebot_plugin_parser/download/__init__.py b/src/nonebot_plugin_parser/download/__init__.py index 9c43ef4e..3de24664 100644 --- a/src/nonebot_plugin_parser/download/__init__.py +++ b/src/nonebot_plugin_parser/download/__init__.py @@ -243,9 +243,9 @@ 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.progress_bar: + with self.progress_bar: + async with aiofiles.open(video_path, "wb") as f: + total_size = 0 update_progress = self.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: From c58e6d8e968d19cbde57e46cc1fbf5755ce2df92 Mon Sep 17 00:00:00 2001 From: fllesser Date: Sun, 2 Aug 2026 15:51:50 +0800 Subject: [PATCH 4/5] tweak --- src/nonebot_plugin_parser/download/__init__.py | 4 +--- tests/parsers/test_acfun.py | 1 + 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/nonebot_plugin_parser/download/__init__.py b/src/nonebot_plugin_parser/download/__init__.py index 3de24664..f153ff6a 100644 --- a/src/nonebot_plugin_parser/download/__init__.py +++ b/src/nonebot_plugin_parser/download/__init__.py @@ -245,14 +245,12 @@ async def download_m3u8( try: with self.progress_bar: async with aiofiles.open(video_path, "wb") as f: - total_size = 0 update_progress = self.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/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: From 00c9f6fecee47c52ae75996a8a6015ce06ba6ace Mon Sep 17 00:00:00 2001 From: fllesser Date: Sun, 2 Aug 2026 16:00:33 +0800 Subject: [PATCH 5/5] tweak --- .../download/__init__.py | 37 ++++--------------- src/nonebot_plugin_parser/download/rich.py | 25 +++++++++++++ 2 files changed, 32 insertions(+), 30 deletions(-) create mode 100644 src/nonebot_plugin_parser/download/rich.py diff --git a/src/nonebot_plugin_parser/download/__init__.py b/src/nonebot_plugin_parser/download/__init__.py index f153ff6a..26734b46 100644 --- a/src/nonebot_plugin_parser/download/__init__.py +++ b/src/nonebot_plugin_parser/download/__init__.py @@ -1,19 +1,13 @@ import asyncio from pathlib import Path -from functools import partial 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 @@ -26,26 +20,9 @@ def __init__(self): self.headers: dict[str, str] = COMMON_HEADER.copy() self.cache_dir: Path = pconfig.cache_dir self.client: httpx.AsyncClient = httpx.AsyncClient(timeout=DOWNLOAD_TIMEOUT, verify=False) - self.progress_bar: Progress = Progress( - TextColumn("[bold blue]{task.description}", justify="right"), - BarColumn(bar_width=None), - "[progress.percentage]{task.percentage:>3.1f}%", - "•", - DownloadColumn(), - ) async def aclose(self): await self.client.aclose() - self.progress_bar.stop() - - def add_progress_task( - self, - desc: str, - total: int | None = None, - ): - task_id = self.progress_bar.add_task(description=desc, total=total) - self.progress_bar.start_task(task_id) - return partial(self.progress_bar.update, task_id) @staticmethod def _validate_content_length( @@ -84,8 +61,8 @@ async def _download_file_with_httpx( response.raise_for_status() content_length = self._validate_content_length(response) - with self.progress_bar: - update_progress = self.add_progress_task( + with progress_bar: + update_progress = add_progress_task( f"httpx | {file_path.name}", content_length, ) @@ -113,8 +90,8 @@ async def _download_file_with_curl_cffi( response.raise_for_status() content_length = self._validate_content_length(response) - with self.progress_bar: - update_progress = self.add_progress_task( + with progress_bar: + update_progress = add_progress_task( f"curl_cffi | {file_path.name}", content_length, ) @@ -243,9 +220,9 @@ async def download_m3u8( video_path = pconfig.cache_dir / video_name try: - with self.progress_bar: + with progress_bar: async with aiofiles.open(video_path, "wb") as f: - update_progress = self.add_progress_task(desc=video_name) + 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): 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)