kmoe-manga-downloader 1.1.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

kmdr/module/downloader/utils.py

@@ -1,142 +1,315 @@
-from typing import Callable, Optional, Union
+import asyncio
 import os
-import time
-from functools import wraps
-
-from requests import Session, HTTPError
-from requests.exceptions import ChunkedEncodingError
-from tqdm import tqdm
 import re
+import math
+from typing import Callable, Optional, Union, Awaitable
+from enum import Enum
+
+from deprecation import deprecated
+import aiohttp
+import aiofiles
+import aiofiles.os as aio_os
+from rich.progress import Progress
+from aiohttp.client_exceptions import ClientPayloadError
 
 BLOCK_SIZE_REDUCTION_FACTOR = 0.75
 MIN_BLOCK_SIZE = 2048
 
-def download_file(
-            session: Session, 
-            url: Union[str, Callable[[], str]],
-            dest_path: str, 
-            filename: str, 
-            retry_times: int = 0, 
-            headers: Optional[dict] = None, 
-            callback: Optional[Callable] = None,
-            block_size: int = 8192
-    ):
+class STATUS(Enum):
+    WAITING='[blue]等待中[/blue]'
+    DOWNLOADING='[cyan]下载中[/cyan]'
+    RETRYING='[yellow]重试中[/yellow]'
+    MERGING='[magenta]合并中[/magenta]'
+    COMPLETED='[green]完成[/green]'
+    FAILED='[red]失败[/red]'
+
+@deprecated(details="请使用 'download_file_multipart'")
+async def download_file(
+        session: aiohttp.ClientSession,
+        semaphore: asyncio.Semaphore,
+        progress: Progress,
+        url: Union[str, Callable[[], str], Callable[[], Awaitable[str]]],
+        dest_path: str,
+        filename: str,
+        retry_times: int = 3,
+        headers: Optional[dict] = None,
+        callback: Optional[Callable] = None,
+):
     """
     下载文件
 
-    :param session: requests.Session 对象
+    :param session: aiohttp.ClientSession 对象
+    :param semaphore: 控制并发的信号量
+    :param progress: 进度条对象
     :param url: 下载链接或者其 Supplier
     :param dest_path: 目标路径
     :param filename: 文件名
     :param retry_times: 重试次数
     :param headers: 请求头
     :param callback: 下载完成后的回调函数
-    :param block_size: 块大小
     """
     if headers is None:
         headers = {}
-    filename_downloading = f'{filename}.downloading'
 
-    file_path = f'{dest_path}/{filename}'
-    tmp_file_path = f'{dest_path}/{filename_downloading}'
-    
-    if not os.path.exists(dest_path):
-        os.makedirs(dest_path, exist_ok=True)
-        
-    if os.path.exists(file_path):
-        tqdm.write(f"{filename} already exists.")
+    file_path = os.path.join(dest_path, filename)
+    filename_downloading = f'{file_path}.downloading'
+
+    if not await aio_os.path.exists(dest_path):
+        await aio_os.makedirs(dest_path, exist_ok=True)
+
+    if await aio_os.path.exists(file_path):
+        progress.console.print(f"[yellow]{filename} 已经存在[/yellow]")
         return
-    
-    if callable(url):
-        url = url()
 
-    resume_from = 0
-    total_size_in_bytes = 0
-    
-    if os.path.exists(tmp_file_path):
-        resume_from = os.path.getsize(tmp_file_path)
-    
-    if resume_from:
-        headers['Range'] = f'bytes={resume_from}-'
+    block_size = 8192
+    attempts_left = retry_times + 1
+    task_id = None
 
     try:
-        with session.get(url = url, stream=True, headers=headers) as r:
-            r.raise_for_status()
+        while attempts_left > 0:
+            attempts_left -= 1
             
-            total_size_in_bytes = int(r.headers.get('content-length', 0)) + resume_from
+            resume_from = (await aio_os.stat(filename_downloading)).st_size if await aio_os.path.exists(filename_downloading) else 0
             
-            with open(tmp_file_path, 'ab') as f:
-                with tqdm(total=total_size_in_bytes, unit='B', unit_scale=True, desc=f'{filename}', initial=resume_from) as progress_bar:
-                    for chunk in r.iter_content(chunk_size=block_size):
-                        if chunk:
-                            f.write(chunk)
-                            progress_bar.update(len(chunk))
+            if resume_from:
+                headers['Range'] = f'bytes={resume_from}-'
+
+            try:
+                async with semaphore:
+                    current_url = await fetch_url(url)
+                    async with session.get(url=current_url, headers=headers) as r:
+                        r.raise_for_status()
+
+                        total_size_in_bytes = int(r.headers.get('content-length', 0)) + resume_from
+
+                        if task_id is None:
+                            task_id = progress.add_task("download", filename=filename, total=total_size_in_bytes, completed=resume_from, status=STATUS.DOWNLOADING.value)
+                        else:
+                            progress.update(task_id, total=total_size_in_bytes, completed=resume_from, status=STATUS.DOWNLOADING.value, refresh=True)
+                        
+                        async with aiofiles.open(filename_downloading, 'ab') as f:
+                            async for chunk in r.content.iter_chunked(block_size):
+                                if chunk:
+                                    await f.write(chunk)
+                                    progress.update(task_id, advance=len(chunk))
+                        
+                        break 
             
-            if (os.path.getsize(tmp_file_path) == total_size_in_bytes):
-                os.rename(tmp_file_path, file_path)
-                
-                if callback:
-                    callback()
+            except Exception as e:
+                if attempts_left > 0:
+                    if task_id is not None:
+                        progress.update(task_id, status=STATUS.RETRYING.value, refresh=True)
+                    if isinstance(e, ClientPayloadError):
+                        new_block_size = max(int(block_size * BLOCK_SIZE_REDUCTION_FACTOR), MIN_BLOCK_SIZE)
+                        if new_block_size < block_size:
+                            block_size = new_block_size
+                    await asyncio.sleep(3)
+                else:
+                    raise e
+        
+        else: 
+            raise IOError(f"Failed to download {filename} after {retry_times} retries.")
+
+        os.rename(filename_downloading, file_path)
+    
     except Exception as e:
-        prefix = f"{type(e).__name__} occurred while downloading {filename}. "
-
-        new_block_size = block_size
-        if isinstance(e, ChunkedEncodingError):
-            new_block_size = max(int(block_size * BLOCK_SIZE_REDUCTION_FACTOR), MIN_BLOCK_SIZE)
-
-        if retry_times > 0:
-            # 重试下载
-            tqdm.write(f"{prefix} Retry after 3 seconds...")
-            time.sleep(3) # 等待3秒后重试，避免触发限流
-            download_file(session, url, dest_path, filename, retry_times - 1, headers, callback, new_block_size)
-        else:
-            tqdm.write(f"{prefix} Meet max retry times, download failed.")
-            raise e
+        if task_id is not None:
+            progress.update(task_id, status=STATUS.FAILED.value, visible=False)
 
-def safe_filename(name: str) -> str:
+    finally:
+        if await aio_os.path.exists(file_path):
+            if task_id is not None:
+                progress.update(task_id, status=STATUS.COMPLETED.value, visible=False)
+
+            if callback:
+                callback()
+
+async def download_file_multipart(
+        session: aiohttp.ClientSession,
+        semaphore: asyncio.Semaphore,
+        progress: Progress,
+        url: Union[str, Callable[[], str], Callable[[], Awaitable[str]]],
+        dest_path: str,
+        filename: str,
+        retry_times: int = 3,
+        chunk_size_mb: int = 10,
+        headers: Optional[dict] = None,
+        callback: Optional[Callable] = None,
+):
     """
-    替换非法文件名字符为下划线
+    下载文件
+
+    :param session: aiohttp.ClientSession 对象
+    :param semaphore: 控制并发的信号量
+    :param progress: 进度条对象
+    :param url: 下载链接或者其 Supplier
+    :param dest_path: 目标路径
+    :param filename: 文件名
+    :param retry_times: 重试次数
+    :param headers: 请求头
+    :param callback: 下载完成后的回调函数
     """
-    return re.sub(r'[\\/:*?"<>|]', '_', name)
+    if headers is None:
+        headers = {}
+        
+    file_path = os.path.join(dest_path, filename)
+    filename_downloading = f'{file_path}.downloading'
+    
+    if not await aio_os.path.exists(dest_path):
+        await aio_os.makedirs(dest_path, exist_ok=True)
+
+    if await aio_os.path.exists(file_path):
+        progress.console.print(f"[blue]{filename} 已经存在[/blue]")
+        return
 
+    part_paths = []
+    task_id = None
+    try:
+        current_url = await fetch_url(url)
 
-function_cache = {}
+        async with session.head(current_url, headers=headers, allow_redirects=True) as response:
+            response.raise_for_status()
+            total_size = int(response.headers['Content-Length'])
 
-def cached_by_kwargs(func):
-    """
-    根据关键字参数缓存函数结果的装饰器。
-
-    Example:
-    >>> @kwargs_cached
-    >>> def add(a, b, c):
-    >>>     return a + b + c
-    >>> result1 = add(1, 2, c=3)  # Calls the function
-    >>> result2 = add(3, 2, c=3)  # Uses cached result
-    >>> assert result1 == result2  # Both results are the same
-    """
+        chunk_size = chunk_size_mb * 1024 * 1024
+        num_chunks = math.ceil(total_size / chunk_size)
+
+        tasks = []
+        
+        resumed_size = 0
+        for i in range(num_chunks):
+            part_path = os.path.join(dest_path, f"{filename}.{i + 1:03d}.downloading")
+            part_paths.append(part_path)
+            if await aio_os.path.exists(part_path):
+                resumed_size += (await aio_os.stat(part_path)).st_size
+
+        task_id = progress.add_task("download", filename=filename, status=STATUS.WAITING.value, total=total_size, completed=resumed_size)
+
+        for i, start in enumerate(range(0, total_size, chunk_size)):
+            end = min(start + chunk_size - 1, total_size - 1)
+
+            task = _download_part(
+                session=session,
+                semaphore=semaphore,
+                url=current_url,
+                start=start,
+                end=end,
+                part_path=part_paths[i],
+                progress=progress,
+                task_id=task_id,
+                headers=headers,
+                retry_times=retry_times
+            )
+            tasks.append(task)
+            
+        await asyncio.gather(*tasks)
+
+        progress.update(task_id, status=STATUS.MERGING.value, refresh=True)
+        await _merge_parts(part_paths, filename_downloading)
+        
+        os.rename(filename_downloading, file_path)
+    except Exception as e:
+        if task_id is not None:
+            progress.update(task_id, status=STATUS.FAILED.value, visible=False)
 
-    global function_cache
-    if func not in function_cache:
-        function_cache[func] = {}
+    finally:
+        if await aio_os.path.exists(file_path):
+            if task_id is not None:
+                progress.update(task_id, status=STATUS.COMPLETED.value, completed=total_size, refresh=True)
 
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        if not kwargs:
-            return func(*args, **kwargs)
+            cleanup_tasks = [aio_os.remove(p) for p in part_paths if await aio_os.path.exists(p)]
+            if cleanup_tasks:
+                await asyncio.gather(*cleanup_tasks)
+            if callback:
+                callback()
 
-        key = frozenset(kwargs.items())
+async def _download_part(
+        session: aiohttp.ClientSession,
+        semaphore: asyncio.Semaphore,
+        url: str,
+        start: int,
+        end: int,
+        part_path: str,
+        progress: Progress,
+        task_id,
+        headers: Optional[dict] = None,
+        retry_times: int = 3
+):
+    if headers is None:
+        headers = {}
+    
+    local_headers = headers.copy()
+    block_size = 8192
+    attempts_left = retry_times + 1
+
+    while attempts_left > 0:
+        attempts_left -= 1
+        
+        try:
+            resume_from = (await aio_os.path.getsize(part_path)) if await aio_os.path.exists(part_path) else 0
+            
+            if resume_from >= (end - start + 1):
+                return
+
+            current_start = start + resume_from
+            local_headers['Range'] = f'bytes={current_start}-{end}'
+            
+            async with semaphore:
+                async with session.get(url, headers=local_headers) as response:
+                    response.raise_for_status()
 
-        if key not in function_cache[func]:
-            function_cache[func][key] = func(*args, **kwargs)
-        return function_cache[func][key]
+                    if progress.tasks[task_id].fields.get("status") != STATUS.DOWNLOADING.value:
+                        progress.update(task_id, status=STATUS.DOWNLOADING.value, refresh=True)
 
-    return wrapper
+                    async with aiofiles.open(part_path, 'ab') as f:
+                        async for chunk in response.content.iter_chunked(block_size):
+                            if chunk:
+                                await f.write(chunk)
+                                progress.update(task_id, advance=len(chunk))
+            return
+        except Exception as e:
+            if attempts_left > 0:
+                await asyncio.sleep(3)
+            else:
+                # console.print(f"[red]分片 {os.path.basename(part_path)} 下载失败: {e}[/red]")
+                raise
+
+async def _merge_parts(part_paths: list[str], final_path: str):
+    async with aiofiles.open(final_path, 'wb') as final_file:
+        try:
+            for part_path in part_paths:
+                async with aiofiles.open(part_path, 'rb') as part_file:
+                    while True:
+                        chunk = await part_file.read(8192)
+                        if not chunk:
+                            break
+                        await final_file.write(chunk)
+        except Exception as e:
+            if aio_os.path.exists(final_path):
+                await aio_os.remove(final_path)
+            raise e
 
-def clear_cache(func):
-    assert hasattr(func, "__wrapped__"), "Function is not wrapped"
-    global function_cache
 
-    wrapped = func.__wrapped__
 
-    if wrapped in function_cache:
-        function_cache[wrapped] = {}
+def safe_filename(name: str) -> str:
+    """
+    替换非法文件名字符为下划线
+    """
+    return re.sub(r'[\\/:*?"<>|]', '_', name)
+
+async def fetch_url(url: Union[str, Callable[[], str], Callable[[], Awaitable[str]]], retry_times: int = 3) -> str:
+    while retry_times >= 0:
+        try:
+            if callable(url):
+                result = url()
+                if asyncio.iscoroutine(result) or isinstance(result, Awaitable):
+                    return await result
+                return result
+            elif isinstance(url, str):
+                return url
+        except Exception as e:
+            retry_times -= 1
+            if retry_times < 0:
+                raise e
+            await asyncio.sleep(2)
+    raise RuntimeError("Max retries exceeded")

kmdr/module/lister/BookUrlLister.py

@@ -10,6 +10,7 @@ class BookUrlLister(Lister):
         super().__init__(*args, **kwargs)
         self._book_url = book_url
 
-    def list(self) -> tuple[BookInfo, list[VolInfo]]:
-        book_info, volumes = extract_book_info_and_volumes(self._session, self._book_url)
-        return book_info, volumes
+    async def list(self) -> tuple[BookInfo, list[VolInfo]]:
+        with self._console.status("获取书籍信息..."):
+            book_info, volumes = await extract_book_info_and_volumes(self._session, self._book_url)
+            return book_info, volumes

kmdr/module/lister/FollowedBookLister.py

@@ -1,6 +1,10 @@
+import asyncio
 from bs4 import BeautifulSoup
+from rich.table import Table
+from rich.prompt import IntPrompt
 
 from kmdr.core import Lister, LISTERS, BookInfo, VolInfo
+from kmdr.core.utils import async_retry
 
 from .utils import extract_book_info_and_volumes
 
@@ -12,27 +16,60 @@ class FollowedBookLister(Lister):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-    def list(self) -> tuple[BookInfo, list[VolInfo]]:
-        followed_rows = BeautifulSoup(self._session.get(url = MY_FOLLOW_URL).text, 'html.parser').find_all('tr', style='height:36px;')
-        mapped = map(lambda x: x.find_all('td'), followed_rows)
-        filtered = filter(lambda x: '書名' not in x[1].text, mapped)
-        books = map(lambda x: BookInfo(name = x[1].text, url = x[1].find('a')['href'], author = x[2].text, status = x[-1].text, last_update = x[-2].text, id = ''), filtered)
-        books = list(books)
+    async def list(self) -> tuple[BookInfo, list[VolInfo]]:
+        books = []
+        
+        with self._console.status("正在获取关注列表..."):
+            books = await self._list_followed_books()
+        
+        if not books:
+            self._console.print("[yellow]关注列表为空。[/yellow]")
+            exit(0)
+
+        table = Table(title="关注的书籍列表", show_header=True, header_style="bold blue")
+        table.add_column("序号", style="dim", width=4, justify="center")
+        table.add_column("书名", style="cyan", no_wrap=True)
+        table.add_column("作者", style="green")
+        table.add_column("最后更新", style="yellow")
+        table.add_column("状态", style="blue")
+
+        for idx, book in enumerate(books):
+            table.add_row(
+                str(idx + 1),
+                book.name,
+                book.author,
+                book.last_update,
+                book.status
+            )
+        
+        self._console.print(table)
+
+        valid_choices = [str(i) for i in range(1, len(books) + 1)]
+        
+        chosen_idx = await asyncio.to_thread(
+            IntPrompt.ask,
+            "请选择要下载的书籍序号",
+            choices=valid_choices,
+            show_choices=False,
+            show_default=False
+        )
+        
+        book_to_download = books[chosen_idx - 1]
+
+        with self._console.status(f"正在获取 '{book_to_download.name}' 的详细信息..."):
+            book_info, volumes = await extract_book_info_and_volumes(self._session, book_to_download.url, book_to_download)
+            return book_info, volumes
+    
+    @async_retry()
+    async def _list_followed_books(self) -> 'list[BookInfo]':
+        async with self._session.get(MY_FOLLOW_URL) as response:
+            response.raise_for_status()
+            html_text = await response.text()
 
-        print("\t最后更新时间\t书名")
-        for v in range(len(books)):
-            print(f"[{v + 1}]\t{books[v].last_update}\t{books[v].name}")
+            # 如果后续有性能问题，可以先考虑使用 lxml 进行解析
+            followed_rows = BeautifulSoup(html_text, 'html.parser').find_all('tr', style='height:36px;')
+            mapped = map(lambda x: x.find_all('td'), followed_rows)
+            filtered = filter(lambda x: '書名' not in x[1].text, mapped)
+            books = list(map(lambda x: BookInfo(name=x[1].text.strip(), url=x[1].find('a')['href'], author=x[2].text.strip(), status=x[-1].text.strip(), last_update=x[-2].text.strip(), id=''), filtered))
         
-        choosed = input("choose a book to download: ")
-        while not choosed.isdigit() or int(choosed) > len(books) or int(choosed) < 1:
-            choosed = input("choose a book to download: ")
-        choosed = int(choosed) - 1
-        book = books[choosed]
-
-        book_info, volumes = extract_book_info_and_volumes(self._session, book.url)
-        book_info.author = book.author
-        book_info.status = book.status
-        book_info.last_update = book.last_update
-
-        return book_info, volumes
-        
+        return books

kmdr/module/lister/utils.py

@@ -1,25 +1,33 @@
-from requests import Session
 from bs4 import BeautifulSoup
 import re
+from typing import Optional
+
+from aiohttp import ClientSession as Session
 
 from kmdr.core import BookInfo, VolInfo, VolumeType
+from kmdr.core.utils import async_retry
 
-def extract_book_info_and_volumes(session: Session, url: str) -> tuple[BookInfo, list[VolInfo]]:
+@async_retry()
+async def extract_book_info_and_volumes(session: Session, url: str, book_info: Optional[BookInfo] = None) -> tuple[BookInfo, list[VolInfo]]:
     """
     从指定的书籍页面 URL 中提取书籍信息和卷信息。
 
-    :param session: 已经建立的 requests.Session 实例。
+    :param session: 已经建立的 HTTP 会话。
     :param url: 书籍页面的 URL。
     :return: 包含书籍信息和卷信息的元组。
     """
-    book_page = BeautifulSoup(session.get(url).text, 'html.parser')
+    async with session.get(url) as response:
+        response.raise_for_status()
+
+        # 如果后续有性能问题，可以先考虑使用 lxml 进行解析
+        book_page = BeautifulSoup(await response.text(), 'html.parser')
 
-    book_info = __extract_book_info(url, book_page)
-    volumes = __extract_volumes(session, book_page)
+        book_info = __extract_book_info(url, book_page, book_info)
+        volumes = await __extract_volumes(session, book_page)
 
-    return book_info, volumes
+        return book_info, volumes
 
-def __extract_book_info(url: str, book_page: BeautifulSoup) -> BookInfo:
+def __extract_book_info(url: str, book_page: BeautifulSoup, book_info: Optional[BookInfo]) -> BookInfo:
     book_name = book_page.find('font', class_='text_bglight_big').text
 
     id = book_page.find('input', attrs={'name': 'bookid'})['value']
@@ -28,35 +36,38 @@ def __extract_book_info(url: str, book_page: BeautifulSoup) -> BookInfo:
         id = id,
         name = book_name,
         url = url,
-        author = '',
-        status = '',
-        last_update = ''
+        author = book_info.author if book_info else '',
+        status = book_info.status if book_info else '',
+        last_update = book_info.last_update if book_info else ''
     )
     
 
-def __extract_volumes(session: Session, book_page: BeautifulSoup) -> list[VolInfo]:
+async def __extract_volumes(session: Session, book_page: BeautifulSoup) -> list[VolInfo]:
     script = book_page.find_all('script', language="javascript")[-1].text
 
     pattern = re.compile(r'/book_data.php\?h=\w+')
     book_data_url = pattern.search(script).group(0)
     
-    book_data = session.get(url = f"https://kox.moe{book_data_url}").text.split('\n')
-    book_data = filter(lambda x: 'volinfo' in x, book_data)
-    book_data = map(lambda x: x.split("\"")[1], book_data)
-    book_data = map(lambda x: x[8:].split(','), book_data)
-    
-    volume_data = list(map(lambda x: VolInfo(
-            id = x[0],
-            extra_info = __extract_extra_info(x[1]),
-            is_last = x[2] == '1',
-            vol_type = __extract_volume_type(x[3]),
-            index = int(x[4]),
-            pages = int(x[6]),
-            name = x[5],
-            size = float(x[11])), book_data))
-    volume_data: list[VolInfo] = volume_data
+    async with session.get(url = f"https://kox.moe{book_data_url}") as response:
+        response.raise_for_status()
+
+        book_data = (await response.text()).split('\n')
+        book_data = filter(lambda x: 'volinfo' in x, book_data)
+        book_data = map(lambda x: x.split("\"")[1], book_data)
+        book_data = map(lambda x: x[8:].split(','), book_data)
+        
+        volume_data = list(map(lambda x: VolInfo(
+                id = x[0],
+                extra_info = __extract_extra_info(x[1]),
+                is_last = x[2] == '1',
+                vol_type = __extract_volume_type(x[3]),
+                index = int(x[4]),
+                pages = int(x[6]),
+                name = x[5],
+                size = float(x[11])), book_data))
+        volume_data: list[VolInfo] = volume_data
 
-    return volume_data
+        return volume_data
 
 def __extract_extra_info(value: str) -> str:
     if value == '0':

kmdr/module/picker/ArgsFilterPicker.py

@@ -27,7 +27,7 @@ class ArgsFilterPicker(Picker):
             volume_data = filter(lambda x: x.index in choice, volume_data)
 
         if self._max_size is not None:
-            volume_data = filter(lambda x: x.size <= self._max_size, volume_data)
+            volume_data = filter(lambda x: self._max_size is None or x.size <= self._max_size, volume_data)
 
         if self._limit is not None:
             return list(volume_data)[:self._limit]