Merge pull request #570 from hadwinfu/fix/download

Fix: 修改视频下载方法为流式下载
This commit is contained in:
Evil0ctal 2025-03-01 13:58:31 -08:00 committed by GitHub
commit 71d6c64490
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -4,7 +4,7 @@ import zipfile
import aiofiles
import httpx
import yaml
from fastapi import APIRouter, Request, Query # 导入FastAPI组件
from fastapi import APIRouter, Request, Query, HTTPException # 导入FastAPI组件
from starlette.responses import FileResponse
from app.api.models.APIResponseModel import ErrorResponseModel # 导入响应模型
@ -18,7 +18,6 @@ config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.pa
with open(config_path, 'r', encoding='utf-8') as file:
config = yaml.safe_load(file)
async def fetch_data(url: str, headers: dict = None):
headers = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
@ -28,6 +27,23 @@ async def fetch_data(url: str, headers: dict = None):
response.raise_for_status() # 确保响应是成功的
return response
# 下载视频专用
async def fetch_data_stream(url: str, headers: dict = None, file_path: str = None):
headers = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
} if headers is None else headers.get('headers')
async with httpx.AsyncClient() as client:
# 启用流式请求
async with client.stream("GET", url, headers=headers) as response:
response.raise_for_status()
# 流式保存文件
async with aiofiles.open(file_path, 'wb') as out_file:
async for chunk in response.aiter_bytes():
await out_file.write(chunk)
return True
@router.get("/download", summary="在线下载抖音|TikTok视频/图片/Online download Douyin|TikTok video/image")
async def download_file_hybrid(request: Request,
@ -104,11 +120,18 @@ async def download_file_hybrid(request: Request,
# 获取视频文件
__headers = await HybridCrawler.TikTokWebCrawler.get_tiktok_headers() if platform == 'tiktok' else await HybridCrawler.DouyinWebCrawler.get_douyin_headers()
response = await fetch_data(url, headers=__headers)
# response = await fetch_data(url, headers=__headers)
# 保存文件
async with aiofiles.open(file_path, 'wb') as out_file:
await out_file.write(response.content)
is_successful = await fetch_data_stream(url, headers=__headers, file_path=file_path)
if not is_successful:
raise HTTPException(
status_code=400,
detail="An error occurred while fetching data"
)
# # 保存文件
# async with aiofiles.open(file_path, 'wb') as out_file:
# await out_file.write(response.content)
# 返回文件内容
return FileResponse(path=file_path, filename=file_name, media_type="video/mp4")