Skip to content

Commit 65e9be7

Browse files
committed
feat(download): 支持分块传输编码的文件下载和实时大小限制
- 解析 Content-Length 头部,支持 None 值(分块传输场景) - 当服务端不提供 Content-Length 时,按实际下载大小进行限流检查 - 优化进度条显示逻辑,处理无总大小的情况 - 增加对空数据块的过滤处理 - 添加更详细的日志记录和异常处理
1 parent 632431a commit 65e9be7

1 file changed

Lines changed: 35 additions & 12 deletions

File tree

src/nonebot_plugin_parser/download/__init__.py

Lines changed: 35 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -67,25 +67,48 @@ async def streamd(
6767
"GET", url, headers=headers, follow_redirects=True
6868
) as response:
6969
response.raise_for_status()
70-
content_length = response.headers.get("Content-Length")
71-
content_length = int(content_length) if content_length else 0
70+
content_length_header = response.headers.get("Content-Length")
7271

73-
if content_length == 0:
74-
logger.warning(f"媒体 url: {url}, 大小为 0, 取消下载")
75-
raise ZeroSizeException
72+
# 解析 Content-Length,允许为 None(分块传输)
73+
content_length = (
74+
int(content_length_header) if content_length_header else None
75+
)
76+
77+
# 有 Content-Length 时,先做一次快速预判
78+
if content_length is not None:
79+
if content_length == 0:
80+
logger.warning(f"媒体 url: {url}, 大小为 0, 取消下载")
81+
raise ZeroSizeException
7682

77-
if (file_size := content_length / 1024 / 1024) > pconfig.max_size:
78-
logger.warning(
79-
f"媒体 url: {url} 大小 {file_size:.2f} MB 超过 {pconfig.max_size} MB, 取消下载"
80-
)
81-
raise SizeLimitException
83+
file_size_mb = content_length / 1024 / 1024
84+
if file_size_mb > pconfig.max_size:
85+
logger.warning(
86+
f"媒体 url: {url} 大小 {file_size_mb:.2f} MB 超过 {pconfig.max_size} MB, 取消下载"
87+
)
88+
raise SizeLimitException
8289

83-
with self.get_progress_bar(file_name, content_length) as bar:
90+
with self.get_progress_bar(file_name, content_length or 0) as bar:
8491
task_id = bar.task_ids[0]
92+
downloaded_bytes = 0
93+
8594
async with aiofiles.open(file_path, "wb") as file:
8695
async for chunk in response.aiter_bytes(1024 * 1024):
96+
if not chunk:
97+
continue
8798
await file.write(chunk)
88-
bar.advance(task_id, len(chunk))
99+
chunk_len = len(chunk)
100+
downloaded_bytes += chunk_len
101+
# 进度条在 content_length 为空时只显示已下载大小
102+
bar.advance(task_id, chunk_len)
103+
104+
# 当服务端不提供 Content-Length 时,按实际已下载大小限流
105+
if content_length is None:
106+
file_size_mb = downloaded_bytes / 1024 / 1024
107+
if file_size_mb > pconfig.max_size:
108+
logger.warning(
109+
f"媒体 url: {url} 实际下载大小 {file_size_mb:.2f} MB 超过 {pconfig.max_size} MB, 取消下载"
110+
)
111+
raise SizeLimitException
89112
# 下载成功,跳出循环
90113
break
91114
except (HTTPError, ConnectionError, TimeoutError, OSError) as e:

0 commit comments

Comments
 (0)