abdds 0.1.0__py3-none-any.whl

abdds/download.py

@@ -0,0 +1,125 @@
+"""下载相关逻辑 - 异步流式下载、断点续传"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, AsyncGenerator
+
+import aiofiles
+import httpx
+
+from ._http import UserAgent
+from .errors import BaiduPanNetworkError
+
+if TYPE_CHECKING:
+    from ._http import _HttpTransport
+
+logger = logging.getLogger("abdds")
+
+
+async def download(
+    transport: _HttpTransport,
+    dlink: str,
+    file_to: Path | None = None,
+    chunk_size: int = 1024 * 1024,
+    max_retries: int = 5,
+) -> AsyncGenerator[bytes, None] | Path:
+    """
+    异步下载文件
+
+    - file_to=None: 返回异步字节迭代器 (流式下载)
+    - file_to=Path: 下载到本地文件，返回 Path
+
+    Args:
+        transport: 异步 HTTP 传输层
+        dlink: 下载链接
+        file_to: 本地保存路径 (Path)，为 None 时返回迭代器
+        chunk_size: 分块大小，默认 1MB
+        max_retries: 最大重试次数，默认 5
+
+    Returns:
+        AsyncGenerator[bytes] 当 file_to=None，Path 当 file_to 指定路径
+
+    Raises:
+        TypeError: file_to 类型不正确
+    """
+    if file_to is not None and not isinstance(file_to, Path):
+        raise TypeError(
+            f"file_to must be Path or None, got {type(file_to).__name__}"
+        )
+    if chunk_size <= 0:
+        raise ValueError(f"chunk_size must be positive, got {chunk_size}")
+
+    stream = _download_stream(transport, dlink, chunk_size, max_retries)
+
+    if file_to is None:
+        return stream
+
+    file_to.parent.mkdir(parents=True, exist_ok=True)
+
+    async with aiofiles.open(file_to, "wb") as f:
+        async for chunk in stream:
+            await f.write(chunk)
+
+    logger.info("Downloaded to %s", file_to)
+    return file_to
+
+
+async def _download_stream(
+    transport: _HttpTransport,
+    dlink: str,
+    chunk_size: int = 1024 * 1024,
+    max_retries: int = 5,
+) -> AsyncGenerator[bytes, None]:
+    """
+    通过 dlink 异步流式下载文件，支持断点续传和指数重试
+    """
+    downloaded_bytes = 0
+    retry_count = 0
+
+    while True:
+        try:
+            headers = {"User-Agent": UserAgent}
+
+            if downloaded_bytes > 0:
+                headers["Range"] = f"bytes={downloaded_bytes}-"
+                logger.info("Resuming download from byte %d...", downloaded_bytes)
+
+            async with transport.client.stream(
+                "GET",
+                dlink,
+                headers=headers,
+                params={"access_token": transport._access_token},
+                timeout=httpx.Timeout(60, connect=15),
+            ) as response:
+                response.raise_for_status()
+
+                # 请求了 Range 但服务器返回 200，说明不支持断点续传
+                if downloaded_bytes > 0 and response.status_code == 200:
+                    logger.warning("Server does not support Range, restarting from beginning")
+                    downloaded_bytes = 0
+
+                async for chunk in response.aiter_bytes(chunk_size=chunk_size):
+                    if chunk:
+                        yield chunk
+                        downloaded_bytes += len(chunk)
+
+            break  # 下载完成
+
+        except httpx.RequestError as e:
+            retry_count += 1
+            if retry_count > max_retries:
+                raise BaiduPanNetworkError(
+                    f"Download failed after {max_retries} retries "
+                    f"({downloaded_bytes} bytes downloaded): {e}",
+                    url=dlink,
+                ) from e
+
+            wait_time = 2 ** (retry_count - 1)
+            logger.warning(
+                "Download interrupted at %d bytes: %s. Retrying in %ds...",
+                downloaded_bytes, e, wait_time,
+            )
+            await asyncio.sleep(wait_time)

abdds/errors.py

@@ -0,0 +1,38 @@
+"""abdds 异常体系"""
+
+
+class BaiduPanError(Exception):
+    """SDK 基础异常"""
+
+
+class BaiduPanNetworkError(BaiduPanError):
+    """网络通信异常"""
+
+    def __init__(
+        self, message: str, *, url: str = "", status_code: int | None = None
+    ):
+        self.url = url
+        self.status_code = status_code
+        super().__init__(message)
+
+
+class BaiduPanAPIError(BaiduPanError):
+    """API 业务异常 (errno != 0)"""
+
+    def __init__(
+        self, errno: int, errmsg: str = "", *, request_id: str = ""
+    ):
+        self.errno = errno
+        self.errmsg = errmsg
+        self.request_id = request_id
+        detail = f"[errno:{errno}] {errmsg}"
+        if request_id:
+            detail += f" (request_id:{request_id})"
+        super().__init__(detail)
+
+
+class TokenExpiredError(BaiduPanError):
+    """Token 过期或无效"""
+
+    def __init__(self, message: str = "Access token expired or invalid"):
+        super().__init__(message)

abdds/models.py

@@ -0,0 +1,58 @@
+"""abdds 数据模型"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class UploadResult:
+    """上传结果"""
+
+    fs_id: int | None
+    md5: str
+    path: str
+    size: int | None
+
+
+@dataclass
+class _ApiPreUploadResponse:
+    """预上传响应（内部模型）"""
+
+    path: str
+    uploadid: str
+    block_list: list[int]
+    block_list_md5: list[str]
+    file_size: int
+
+
+@dataclass
+class ApiFileMeta:
+    """文件元信息"""
+
+    md5: str
+    path: str
+    size: int
+    fs_id: int
+    dlink: str
+    filename: str
+
+
+@dataclass
+class ApiFileListItem:
+    """文件列表项"""
+
+    md5: str
+    path: str
+    size: int
+    fs_id: int
+    is_dir: int
+    filename: str
+
+
+@dataclass
+class ApiQuotaInfo:
+    """空间配额信息"""
+
+    total: int
+    used: int

abdds/upload.py

@@ -0,0 +1,368 @@
+"""上传相关逻辑 - 异步小文件单步上传 & 大文件分片上传"""
+
+from __future__ import annotations
+
+import hashlib
+import inspect
+import io
+import json
+import logging
+from collections.abc import AsyncGenerator as _AsyncGeneratorABC
+from collections.abc import Generator as _GeneratorABC
+from pathlib import Path, PurePosixPath
+from typing import TYPE_CHECKING, AsyncGenerator, Generator
+
+import aiofiles
+import aiofiles.os
+
+from .errors import BaiduPanAPIError
+from .models import _ApiPreUploadResponse, UploadResult
+
+if TYPE_CHECKING:
+    from ._http import _HttpTransport
+
+logger = logging.getLogger("abdds")
+
+BLOCK_SIZE = 4 * 1024 * 1024  # 4MB
+
+# file_from 参数可接受类型
+UploadSource = Path | bytes | Generator[bytes, None, None] | AsyncGenerator[bytes, None]
+
+
+def _is_generator(obj: object) -> bool:
+    """判断对象是否为同步生成器"""
+    return isinstance(obj, _GeneratorABC) and not isinstance(obj, _AsyncGeneratorABC)
+
+
+def _is_async_generator(obj: object) -> bool:
+    """判断对象是否为异步生成器"""
+    return isinstance(obj, _AsyncGeneratorABC)
+
+
+async def upload(
+    transport: _HttpTransport,
+    file_from: UploadSource,
+    to: PurePosixPath,
+    pan_dir_base: str,
+) -> UploadResult:
+    """
+    异步上传数据到百度网盘
+
+    - 小文件 (<=4MB): 单步上传 (precreate + direct upload)
+    - 大文件 (>4MB):  分片上传 (precreate + 分片 + merge)
+
+    Args:
+        transport: 异步 HTTP 传输层
+        file_from: 数据来源，支持 Path / bytes / Generator[bytes] / AsyncGenerator[bytes]
+        to: 远程路径 (PurePosixPath)，将拼接在 /apps/{app_name} 之后
+        pan_dir_base: 应用目录前缀，如 /apps/myapp
+
+    Raises:
+        TypeError: file_from 类型不正确
+    """
+    # 类型校验
+    if (
+        not isinstance(file_from, (Path, bytes))
+        and not _is_generator(file_from)
+        and not _is_async_generator(file_from)
+    ):
+        raise TypeError(
+            f"file_from must be Path, bytes, Generator[bytes], or AsyncGenerator[bytes], "
+            f"got {type(file_from).__name__}"
+        )
+
+    # 构造完整远程路径: pan_dir_base + to
+    remote_path = f"{pan_dir_base}{to.as_posix()}"
+
+    # 根据 file_from 类型提取数据
+    if isinstance(file_from, Path):
+        if not file_from.is_file():
+            raise FileNotFoundError(f"File not found: {file_from}")
+        file_size = (await aiofiles.os.stat(file_from)).st_size
+        if file_size == 0:
+            raise ValueError("Cannot upload empty file (0 bytes)")
+        block_md5_list = await _calculate_file_md5_blocks(file_from)
+    elif isinstance(file_from, bytes):
+        file_size = len(file_from)
+        if file_size == 0:
+            raise ValueError("Cannot upload empty data (0 bytes)")
+        block_md5_list = _calculate_bytes_md5_blocks(file_from)
+    elif _is_async_generator(file_from):
+        # 异步生成器需要消费一次来计算大小和 MD5，缓存在内存中
+        chunks: list[bytes] = []
+        file_size = 0
+        md5_list: list[str] = []
+        current_chunk = io.BytesIO()
+
+        async for data in file_from:
+            current_chunk.write(data)
+            file_size += len(data)
+            if current_chunk.tell() >= BLOCK_SIZE:
+                block_data = current_chunk.getvalue()
+                chunks.append(block_data)
+                md5_list.append(hashlib.md5(block_data).hexdigest())
+                current_chunk = io.BytesIO()
+
+        # 处理最后不足一个 BLOCK_SIZE 的数据
+        remaining = current_chunk.getvalue()
+        if remaining:
+            chunks.append(remaining)
+            md5_list.append(hashlib.md5(remaining).hexdigest())
+
+        block_md5_list = md5_list
+
+        if file_size == 0:
+            raise ValueError("Cannot upload empty async generator (0 bytes yielded)")
+
+        # 将 chunks 重新组装为异步生成器供上传使用
+        async def _replay_chunks() -> AsyncGenerator[bytes, None]:
+            for chunk in chunks:
+                yield chunk
+
+        file_from = _replay_chunks()
+    elif _is_generator(file_from):
+        # 同步生成器
+        chunks2: list[bytes] = []
+        file_size = 0
+        md5_list2: list[str] = []
+        current_chunk2 = io.BytesIO()
+
+        for data in file_from:
+            current_chunk2.write(data)
+            file_size += len(data)
+            if current_chunk2.tell() >= BLOCK_SIZE:
+                block_data = current_chunk2.getvalue()
+                chunks2.append(block_data)
+                md5_list2.append(hashlib.md5(block_data).hexdigest())
+                current_chunk2 = io.BytesIO()
+
+        remaining2 = current_chunk2.getvalue()
+        if remaining2:
+            chunks2.append(remaining2)
+            md5_list2.append(hashlib.md5(remaining2).hexdigest())
+
+        block_md5_list = md5_list2
+
+        if file_size == 0:
+            raise ValueError("Cannot upload empty generator (0 bytes yielded)")
+
+        file_from = (chunk for chunk in chunks2)
+    else:
+        raise TypeError(
+            f"file_from must be Path, bytes, Generator[bytes], or AsyncGenerator[bytes], "
+            f"got {type(file_from).__name__}"
+        )
+
+    logger.info("Uploading -> %s (%d bytes)", remote_path, file_size)
+
+    # 预上传
+    pre_resp = await _api_pre_create(transport, remote_path, file_size, block_md5_list)
+
+    # 获取上传服务器
+    upload_host = await _get_upload_server(transport, remote_path, pre_resp.uploadid)
+
+    if file_size <= BLOCK_SIZE:
+        return await _upload_small(transport, upload_host, pre_resp, file_from, file_size)
+    else:
+        return await _upload_multipart(transport, upload_host, pre_resp, file_from, block_md5_list)
+
+
+async def _upload_small(
+    transport: _HttpTransport,
+    host: str,
+    pre_resp: _ApiPreUploadResponse,
+    file_from: UploadSource,
+    file_size: int,
+) -> UploadResult:
+    """小文件单步上传"""
+    url = f"{host}/rest/2.0/pcs/file"
+    params = {"method": "upload", "path": pre_resp.path}
+
+    if isinstance(file_from, Path):
+        async with aiofiles.open(file_from, "rb") as f:
+            data = await f.read()
+    elif isinstance(file_from, bytes):
+        data = file_from
+    elif _is_async_generator(file_from):
+        chunks: list[bytes] = []
+        async for chunk in file_from:
+            chunks.append(chunk)
+        data = b"".join(chunks)
+    elif _is_generator(file_from):
+        data = b"".join(file_from)
+    else:
+        data = b""
+
+    files = {"file": ("blob", data)}
+    resp = await transport.request("POST", url, params=params, files=files)
+    return UploadResult(
+        fs_id=resp.get("fs_id"),
+        md5=resp.get("md5"),
+        path=resp.get("path"),
+        size=resp.get("size"),
+    )
+
+
+async def _upload_multipart(
+    transport: _HttpTransport,
+    host: str,
+    pre_resp: _ApiPreUploadResponse,
+    file_from: UploadSource,
+    block_md5_list: list[str],
+) -> UploadResult:
+    """大文件分片上传"""
+    url_prefix = f"{host}/rest/2.0/pcs/superfile2"
+
+    if isinstance(file_from, Path):
+        async with aiofiles.open(file_from, "rb") as f:
+            for idx in range(len(block_md5_list)):
+                await f.seek(idx * BLOCK_SIZE)
+                chunk_data = await f.read(BLOCK_SIZE)
+                await _upload_part(transport, url_prefix, pre_resp, idx, chunk_data)
+    elif isinstance(file_from, bytes):
+        for idx in range(len(block_md5_list)):
+            chunk_data = file_from[idx * BLOCK_SIZE : (idx + 1) * BLOCK_SIZE]
+            await _upload_part(transport, url_prefix, pre_resp, idx, chunk_data)
+    elif _is_async_generator(file_from):
+        idx = 0
+        async for chunk_data in file_from:
+            await _upload_part(transport, url_prefix, pre_resp, idx, chunk_data)
+            idx += 1
+    elif _is_generator(file_from):
+        for idx, chunk_data in enumerate(file_from):
+            await _upload_part(transport, url_prefix, pre_resp, idx, chunk_data)
+
+    # 合并分片
+    return await _api_create_file(
+        transport, pre_resp.path, pre_resp.uploadid, pre_resp.file_size, block_md5_list
+    )
+
+
+async def _upload_part(
+    transport: _HttpTransport,
+    url_prefix: str,
+    pre_resp: _ApiPreUploadResponse,
+    partseq: int,
+    chunk_data: bytes,
+) -> None:
+    """上传单个分片"""
+    logger.debug("Uploading part %d", partseq + 1)
+    params = {
+        "method": "upload",
+        "type": "tmpfile",
+        "path": pre_resp.path,
+        "uploadid": pre_resp.uploadid,
+        "partseq": str(partseq),
+    }
+    files = {"file": ("blob", chunk_data)}
+    await transport.request("POST", url_prefix, params=params, files=files)
+
+
+async def _api_pre_create(
+    transport: _HttpTransport,
+    path: str,
+    size: int,
+    block_list: list[str],
+) -> _ApiPreUploadResponse:
+    """预上传 - 通知云端新建上传任务"""
+    params = {"method": "precreate"}
+    data = {
+        "path": path,
+        "size": size,
+        "isdir": "0",
+        "autoinit": "1",
+        "rtype": "3",  # 覆盖
+        "block_list": json.dumps(block_list),
+    }
+    resp = await transport.request("POST", "https://pan.baidu.com/rest/2.0/xpan/file", params=params, data=data)
+    return _ApiPreUploadResponse(
+        path=path,
+        uploadid=resp["uploadid"],
+        block_list=resp.get("block_list", []),
+        block_list_md5=block_list,
+        file_size=size,
+    )
+
+
+async def _get_upload_server(transport: _HttpTransport, path: str, uploadid: str) -> str:
+    """获取上传服务器域名"""
+    params = {
+        "method": "locateupload",
+        "appid": "250528",
+        "path": path,
+        "uploadid": uploadid,
+        "upload_version": "2.0",
+    }
+    resp = await transport.request("GET", "https://pan.baidu.com/rest/2.0/xpan/file", params=params)
+    servers = resp.get("servers", [])
+    if not servers:
+        raise BaiduPanAPIError(-1, "No upload server available")
+
+    for s in servers:
+        server_url = s.get("server", "")
+        if server_url.startswith("https"):
+            return server_url
+
+    # 降级：返回第一个有效 server
+    for s in servers:
+        server_url = s.get("server", "")
+        if server_url:
+            return server_url
+
+    raise BaiduPanAPIError(-1, "No valid upload server in response")
+
+
+async def _api_create_file(
+    transport: _HttpTransport,
+    path: str,
+    uploadid: str,
+    size: int,
+    block_list: list[str],
+) -> UploadResult:
+    """合并分片完成上传"""
+    params = {"method": "create"}
+    data = {
+        "path": path,
+        "size": size,
+        "isdir": "0",
+        "rtype": "3",
+        "uploadid": uploadid,
+        "block_list": json.dumps(block_list),
+    }
+    resp = await transport.request("POST", "https://pan.baidu.com/rest/2.0/xpan/file", params=params, data=data)
+    return UploadResult(
+        fs_id=resp.get("fs_id"),
+        md5=resp.get("md5"),
+        path=resp.get("path"),
+        size=resp.get("size"),
+    )
+
+
+async def _calculate_file_md5_blocks(file_path: Path) -> list[str]:
+    """异步计算文件分片 MD5"""
+    file_size = (await aiofiles.os.stat(file_path)).st_size
+
+    async with aiofiles.open(file_path, "rb") as f:
+        if file_size <= BLOCK_SIZE:
+            content = await f.read()
+            return [hashlib.md5(content).hexdigest()]
+
+        md5_list: list[str] = []
+        while True:
+            chunk = await f.read(BLOCK_SIZE)
+            if not chunk:
+                break
+            md5_list.append(hashlib.md5(chunk).hexdigest())
+
+    return md5_list
+
+
+def _calculate_bytes_md5_blocks(data: bytes) -> list[str]:
+    """计算 bytes 对象的分片 MD5"""
+    if len(data) <= BLOCK_SIZE:
+        return [hashlib.md5(data).hexdigest()]
+
+    md5_list: list[str] = []
+    for i in range(0, len(data), BLOCK_SIZE):
+        md5_list.append(hashlib.md5(data[i : i + BLOCK_SIZE]).hexdigest())
+    return md5_list