pdman 0.2.0__py3-none-any.whl

pdman/__init__.py

@@ -0,0 +1,4 @@
+from .manager import Manager
+
+manager = Manager()
+__all__ = ["Manager", "manager"]

pdman/chunk.py

@@ -0,0 +1,198 @@
+import random
+import os
+import time
+import asyncio
+import aiohttp
+import aiofiles
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .downloader import Downloader
+
+
+class Chunk:
+    def __init__(
+        self,
+        parent: Downloader,
+        start: int,
+        end: int,
+        chunk_path: str,
+        forward: Chunk = None,
+        next: Chunk = None,
+    ):
+        self.parent = parent
+        self.start = start
+        self.end = end
+        self.chunk_path = chunk_path
+        if os.path.exists(chunk_path):
+            self.size = os.path.getsize(chunk_path)
+        else:
+            self.size = 0
+        self.forward: Chunk = forward
+        self.next: Chunk = next
+
+    def __iter__(self):
+        current = self
+        while current:
+            yield current
+            current = current.next
+
+    def __str__(self):
+        return f"Chunk(start={self.start}, end={self.end},target size={(self.end - self.start + 1) if self.end is not None else -1}, size={self.size}, chunk_path={self.chunk_path})"
+
+    def __add__(self, other):
+        if not isinstance(other, Chunk):
+            return NotImplemented
+        return self.size + other.size
+
+    def __radd__(self, other):
+        if other == 0:
+            return self.size
+        if not isinstance(other, int):
+            return NotImplemented
+        return self.size + other
+
+    def _is_complete(self) -> bool:
+        if self.end is None:
+            self.parent._downloaded = True
+            return self.size > 0
+        return self.end is not None and self.size == self.end - self.start + 1
+
+    def _needs_download(self) -> bool:
+        return self.end is None or self.size < self.end - self.start + 1
+
+    def _apply_range_header(self, headers: dict):
+        if self.end is not None:
+            if self.start + self.size <= self.end:
+                headers["Range"] = f"bytes={self.start + self.size}-{self.end}"
+            else:
+                headers["break"] = True
+        else:
+            if "Range" in headers:
+                headers.pop("Range")
+
+    async def _stream_response(self, response, f) -> bool:
+        last_time = time.time()
+        pos = await f.tell()
+        continue_flag = False
+        async for data in response.content.iter_chunked(10240):
+            if self.end is not None:
+                remaining = self.end - self.start + 1 - pos
+                if remaining <= 0:
+                    break
+                data = data[:remaining]
+            await f.write(data)
+            async with self.parent.lock:
+                self.size += len(data)
+                now = time.time()
+                elaps = max(now - last_time, 1e-6)
+                speed = len(data) / elaps
+                if (
+                    self.parent.parent.chunk_retry_speed
+                    and speed < self.parent.parent.chunk_retry_speed
+                ):
+                    continue_flag = True
+                last_time = now
+            pos += len(data)
+        return continue_flag
+
+    async def _split_incomplete(self):
+        if self.size != self.end - self.start + 1:
+            self.parent._logger.debug(
+                f"Chunk not fully downloaded, splitting chunk: {self}"
+            )
+            async with self.parent.lock:
+                new_start = self.start + self.size
+                new_chunk = Chunk(
+                    self.parent,
+                    new_start,
+                    self.end,
+                    os.path.join(
+                        self.parent.pdm_tmp,
+                        f"{self.parent.filename}.{new_start}",
+                    ),
+                    self,
+                    next=self.next,
+                )
+                self.end = new_start - 1
+                self.next = new_chunk
+
+    async def download(self):
+        assert self.end is not None or self.size >= 0
+        headers = {}  # TODO 添加其他必要的headers
+        file_mode = "ab" if os.path.exists(self.chunk_path) else "wb"
+        async with (
+            aiohttp.ClientSession(
+                timeout=aiohttp.ClientTimeout(sock_read=30)
+            ) as session,
+            aiofiles.open(self.chunk_path, file_mode) as f,
+        ):
+            for _ in range(self.parent.parent.retry):
+                if os.path.exists(self.chunk_path) and self._is_complete():
+                    return self
+                while True:
+                    try:
+                        self._apply_range_header(headers)
+                        if headers.get("break"):
+                            break
+                        self.parent._logger.debug(
+                            f"Downloading chunk: {self}, with headers: {headers}"
+                        )
+                        if self._needs_download():
+                            async with session.get(
+                                self.parent.url,
+                                headers=headers,
+                                timeout=self.parent.parent.chunk_timeout,
+                            ) as response:
+                                if response.status in (200, 206):
+                                    continue_flag = await self._stream_response(
+                                        response, f
+                                    )
+                                    if continue_flag:
+                                        await asyncio.sleep(
+                                            self.parent.parent.retry_wait
+                                        )
+                                        self.parent._logger.debug(
+                                            "speed is low restarting..."
+                                        )
+                                        continue
+                    except aiohttp.client_exceptions.ClientPayloadError:
+                        await asyncio.sleep(
+                            self.parent.parent.retry_wait + random.random() * 5
+                        )
+                    except asyncio.TimeoutError:
+                        self.parent._logger.debug(
+                            f"Timeout downloading chunk {self}, retrying..."
+                        )
+                        await asyncio.sleep(
+                            self.parent.parent.retry_wait + random.random() * 5
+                        )
+                    except ConnectionResetError:
+                        self.parent._logger.debug(
+                            f"Connection reset downloading chunk {self}, retrying..."
+                        )
+                        await asyncio.sleep(
+                            self.parent.parent.retry_wait + random.random() * 5
+                        )
+                    except Exception as e:
+                        self.parent._logger.debug(
+                            f"Error downloading chunk {self}: {e}"
+                        )
+                        # traceback.print_exc()
+                        await asyncio.sleep(self.parent.parent.retry_wait)
+                        break
+                    if self._is_complete():
+                        return self
+                if self.end is None:
+                    self.parent._logger.debug(
+                        f"completed download chunk (unknown size): {self}"
+                    )
+                    self.parent._downloaded = True
+                    return self
+                elif self._needs_download():
+                    self.parent._logger.debug(f"retrying download chunk: {self}")
+                else:
+                    self.parent._logger.debug(f"completed download chunk: {self}")
+        await self._split_incomplete()
+        return self

pdman/cli.py

@@ -0,0 +1,207 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+"""
+CLI 入口：提供命令行解析并调用 PDManager。
+安装后可通过 console_scripts 生成可执行命令。
+"""
+
+import os
+import sys
+import argparse
+import asyncio
+from .manager import Manager
+
+version = "0.2.0"
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-v",
+        "--version",
+        action="version",
+        version=f"PythonDownloadManager(PDMAN) version {version}",
+        help="Print the version number and exit.",
+    )
+    parser.add_argument(
+        "-l",
+        "--log",
+        type=str,
+        required=False,
+        default=None,
+        help="The file name of the log file. If '-' is specified, log is written to stdout.",
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Enable debug mode with verbose logging.",
+    )
+    parser.add_argument(
+        "-d",
+        "--dir",
+        type=str,
+        default=os.path.join(os.getcwd(), "pdm"),
+        help="The directory to store the downloaded file.",
+    )
+    parser.add_argument(
+        "-o",
+        "--out",
+        type=str,
+        default=None,
+        help="The file name of the downloaded file. It is always relative to the directory given in -d option. When the -Z option is used, this option will be ignored.",
+    )
+    parser.add_argument(
+        "-V",
+        "--check-integrity",
+        action="store_true",
+        help="Check file integrity by validating piece hashes or a hash of the entire file.",
+    )
+    parser.add_argument(
+        "-c",
+        "--continue",
+        dest="continue_download",
+        action="store_true",
+        help="Continue downloading a partially downloaded file.",
+    )
+    parser.add_argument(
+        "-i",
+        "--input-file",
+        type=str,
+        default=[],
+        action="append",
+        help="Downloads URIs found in FILE(s). Supports JSON, YAML, or plain text.",
+    )
+    parser.add_argument(
+        "-x",
+        "--max-concurrent-downloads",
+        type=int,
+        default=5,
+        help="Set maximum number of parallel downloads for each URL or task.",
+    )
+    parser.add_argument(
+        "--chunk-retry-speed",
+        default="",
+        help="If the chunk speed falls below SIZE bytes/second, restart that chunk. Append K/M.",
+    )
+    parser.add_argument(
+        "-r",
+        "--retry",
+        type=int,
+        default=3,
+        help="Number of times to retry downloading a URL upon failure.",
+    )
+    parser.add_argument(
+        "-W",
+        "--retry-wait",
+        type=int,
+        default=5,
+        help="Maximum wait time in seconds between retries.",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=int,
+        default=None,
+        help="Timeout in seconds for each download request.",
+    )
+    parser.add_argument(
+        "--chunk-timeout",
+        type=int,
+        default=None,
+        help="Timeout in seconds for each chunk download request.",
+    )
+    parser.add_argument(
+        "-N",
+        "--max-downloads",
+        type=int,
+        default=4,
+        help="The maximum number of concurrent downloads.",
+    )
+    parser.add_argument(
+        "--no-auto-file-renaming",
+        action="store_false",
+        help="Disable auto renaming when target file exists.",
+    )
+    parser.add_argument(
+        "-Z",
+        "--force-sequential",
+        action="store_true",
+        help="Fetch URIs sequentially.",
+    )
+    parser.add_argument(
+        "-k",
+        "--min-split-size",
+        type=str,
+        default="1M",
+        help="Minimum split size. Append K/M.",
+    )
+    parser.add_argument(
+        "--tmp",
+        type=str,
+        default=None,
+        help="Temporary directory for chunk files.",
+    )
+    parser.add_argument(
+        "-t",
+        "--threads",
+        type=int,
+        default=4,
+        help="Alias of max-downloads (deprecated).",
+    )
+    parser.add_argument(
+        "-ua",
+        "--user-agent",
+        type=str,
+        default="PDMAN-Downloader/1.0",
+        help="The User-Agent string to use for HTTP requests.",
+    )
+    parser.add_argument(
+        "urls",
+        type=str,
+        nargs="*",
+        default=None,
+        help="The URL(s) to download.",
+    )
+
+    args = parser.parse_args(argv)
+    if args.log == "-":
+        args.log = sys.stdout
+    if args.force_sequential and args.out is not None:
+        args.out = None
+
+    pdm = Manager(
+        max_downloads=args.max_downloads,
+        log_path=args.log,
+        debug=args.debug,
+        continue_download=args.continue_download,
+        max_concurrent_downloads=args.max_concurrent_downloads,
+        min_split_size=args.min_split_size,
+        force_sequential=args.force_sequential,
+        tmp_dir=args.tmp,
+        check_integrity=args.check_integrity,
+        user_agent=args.user_agent,
+        chunk_retry_speed=args.chunk_retry_speed,
+        retry=args.retry,
+        retry_wait=args.retry_wait,
+        timeout=args.timeout,
+        chunk_timeout=args.chunk_timeout,
+        auto_file_renaming=args.no_auto_file_renaming,
+        out_dir=args.dir,
+    )
+
+    if args.urls and len(args.urls) == 1 and args.out is not None:
+        pdm.append(args.urls[0], file_name=args.out)
+    else:
+        if args.out is not None:
+            pass  # ignore --out when multiple urls
+        pdm.add_urls(args.urls or [])
+
+    if args.input_file:
+        for file in args.input_file:
+            if os.path.exists(file):
+                pdm.load_input_file(file)
+
+    asyncio.run(pdm.download())
+
+
+if __name__ == "__main__":
+    main()