fix(download): fixed progress bar display with rich

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 视频下载失败")

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)

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: