101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
from typing import Optional, Any, Union
|
||
from functools import wraps
|
||
import backoff
|
||
import httpx
|
||
import asyncio
|
||
from loguru import logger
|
||
import aiofiles
|
||
from pathlib import Path
|
||
|
||
|
||
class HTTPDownloadUtils:
|
||
# 创建一个类级别的 Semaphore 来控制并发
|
||
_semaphore = asyncio.Semaphore(5) # 限制最大并发数为5
|
||
|
||
@staticmethod
|
||
@backoff.on_exception(
|
||
backoff.expo,
|
||
(httpx.RequestError, httpx.HTTPStatusError, Exception),
|
||
max_tries=10,
|
||
max_time=300, # 最大重试时间5分钟
|
||
giveup=lambda e: isinstance(e, httpx.HTTPStatusError) and e.response.status_code == 404
|
||
)
|
||
async def async_download_file(
|
||
url: str,
|
||
output_path: Union[str, Path],
|
||
timeout: float = 60.0,
|
||
verify: bool = True
|
||
) -> Union[str, None]:
|
||
"""
|
||
异步下载文件或获取内容
|
||
|
||
:param url: 要下载的URL
|
||
:param output_path: 输出文件路径，如果为None则返回内容
|
||
:param timeout: 请求超时时间（秒）
|
||
:param verify: 是否验证SSL证书
|
||
|
||
:return: 如果output_path为None，返回下载的内容；否则返回保存的文件路径
|
||
|
||
:exception httpx.RequestError: 请求错误
|
||
:exception httpx.HTTPStatusError: HTTP状态错误
|
||
:exception Exception: 其他错误
|
||
"""
|
||
async with HTTPDownloadUtils._semaphore: # 使用信号量控制并发
|
||
try:
|
||
logger.info(f"Starting download from {url}")
|
||
|
||
async with httpx.AsyncClient(
|
||
timeout=timeout,
|
||
verify=verify,
|
||
follow_redirects=True
|
||
) as client:
|
||
response = await client.get(url)
|
||
response.raise_for_status() # 检查HTTP状态码
|
||
|
||
output_path = Path(output_path)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
async with aiofiles.open(output_path, 'wb') as f:
|
||
await f.write(response.content)
|
||
logger.success(f"Successfully downloaded to {output_path}")
|
||
return str(output_path)
|
||
|
||
|
||
except httpx.HTTPStatusError as e:
|
||
logger.error(f"HTTP error occurred: {e.response.status_code} - {url}")
|
||
raise
|
||
except httpx.RequestError as e:
|
||
logger.error(f"Request error occurred: {str(e)} - {url}")
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Unexpected error occurred: {str(e)} - {url}")
|
||
raise
|
||
|
||
@staticmethod
|
||
async def batch_download(
|
||
urls: list[str],
|
||
output_paths: list[Union[str, Path]] = None,
|
||
max_concurrent: int = 5
|
||
) -> list[Union[str, None]]:
|
||
"""
|
||
批量下载多个文件
|
||
|
||
:param urls: URL列表
|
||
:param output_paths: 输出路径列表
|
||
:param max_concurrent: 最大并发数
|
||
|
||
:return: 下载结果列表
|
||
"""
|
||
# 更新信号量的值
|
||
HTTPDownloadUtils._semaphore = asyncio.Semaphore(max_concurrent)
|
||
|
||
tasks = []
|
||
for i, url in enumerate(urls):
|
||
output_path = output_paths[i]
|
||
task = asyncio.create_task(
|
||
HTTPDownloadUtils.async_download_file(url, output_path)
|
||
)
|
||
tasks.append(task)
|
||
|
||
return await asyncio.gather(*tasks, return_exceptions=True)
|