206 lines
7.9 KiB
Python
206 lines
7.9 KiB
Python
"""导出采集图片：把勾选的采集图/生成图打包成 ZIP 下载到本地。
|
||
|
||
平移自 image-suite-studio；路径按 docs/v2.1/api.md §7 定为 POST /api/export/images
|
||
（前端 services/suite.ts exportImages 同款契约，注意与 ISS 的 /export-images 不同）。
|
||
ZIP 内部按分组名建子文件夹（主图 / SKU图片 / 详情图 / 手动上传），
|
||
文件名沿用采集 key（main-001 等）+ SKU 规格名；顶层文件夹用商品标题（清洗后）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import mimetypes
|
||
import re
|
||
import tempfile
|
||
import zipfile
|
||
from collections.abc import Iterator
|
||
from urllib.parse import quote
|
||
|
||
from fastapi import APIRouter, HTTPException
|
||
from fastapi.responses import Response, StreamingResponse
|
||
from pydantic import BaseModel, Field
|
||
|
||
from api.proxy import guess_referer
|
||
from config import get_settings
|
||
from services.storage import (
|
||
ResourceDownloadTimeoutError,
|
||
ResourceSizeLimitError,
|
||
UnsafeRemoteUrlError,
|
||
download_bytes,
|
||
local_path,
|
||
read_local_bytes,
|
||
)
|
||
|
||
router = APIRouter(prefix="/api", tags=["export"])
|
||
|
||
_EXT_BY_CTYPE = {
|
||
"image/jpeg": ".jpg",
|
||
"image/png": ".png",
|
||
"image/webp": ".webp",
|
||
"image/gif": ".gif",
|
||
"image/bmp": ".bmp",
|
||
}
|
||
|
||
|
||
class ExportImageItem(BaseModel):
|
||
url: str
|
||
groupName: str = "主图"
|
||
variantName: str | None = None
|
||
key: str = "" # 采集 key，如 main-001 / sku-002 / upload-001
|
||
|
||
|
||
class ExportImagesRequest(BaseModel):
|
||
title: str | None = Field(default=None, description="商品标题，用作 ZIP 顶层文件夹名")
|
||
images: list[ExportImageItem] = Field(min_length=1)
|
||
|
||
|
||
def _clean(name: str) -> str:
|
||
"""清洗文件夹/文件名非法字符（与插件 cleanFilename 同规则，Windows 兼容）。"""
|
||
s = re.sub(r'[<>:"/\\|?*\x00-\x1f]', "_", (name or "").strip())
|
||
s = re.sub(r"\s+", "_", s)
|
||
return s.strip(" .")[:80]
|
||
|
||
|
||
def _ext(url: str, ctype: str) -> str:
|
||
"""由 content-type（优先）或 URL 后缀决定扩展名。"""
|
||
ctype = ctype.split(";")[0].strip().lower()
|
||
if ctype in _EXT_BY_CTYPE:
|
||
return _EXT_BY_CTYPE[ctype]
|
||
if ctype.startswith("image/"):
|
||
return "." + ctype.split("/")[-1]
|
||
path = url.split("?")[0].lower()
|
||
for ext in (".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp"):
|
||
if path.endswith(ext):
|
||
return ".jpg" if ext == ".jpeg" else ext
|
||
return ".jpg"
|
||
|
||
|
||
def _is_image(url: str, ctype: str) -> bool:
|
||
ctype = ctype.split(";")[0].strip().lower()
|
||
if ctype.startswith("image/"):
|
||
return True
|
||
return bool(re.search(r"\.(jpe?g|png|webp|gif|bmp)(\?|$)", url, re.IGNORECASE))
|
||
|
||
|
||
def _download_disposition(filename: str) -> str:
|
||
"""生成兼容中文文件名的下载响应头。
|
||
|
||
Starlette 会按 Latin-1 编码 HTTP 头；直接将中文放进 filename 会抛 UnicodeEncodeError
|
||
并返回 500。filename 提供 ASCII 回退，filename* 使用 RFC 5987 的 UTF-8 百分号编码。
|
||
"""
|
||
fallback = re.sub(r"[^\x20-\x7e]", "_", filename) or "image"
|
||
return f"attachment; filename=\"{fallback}\"; filename*=UTF-8''{quote(filename, safe='')}"
|
||
|
||
|
||
async def _download(url: str) -> tuple[bytes, str]:
|
||
"""本地 media 文件直读磁盘；远程 URL（源站 CDN / 七牛）带 Referer 下载。"""
|
||
path = local_path(url)
|
||
if path is not None:
|
||
mime = mimetypes.guess_type(path.name)[0] or "image/jpeg"
|
||
return read_local_bytes(path), mime
|
||
return await download_bytes(url, referer=guess_referer(url))
|
||
|
||
|
||
def _stream_file(handle: io.BufferedRandom, chunk_size: int = 64 * 1024) -> Iterator[bytes]:
|
||
"""以固定块流式返回临时 ZIP，响应结束（或中断）后关闭临时文件。"""
|
||
try:
|
||
while chunk := handle.read(chunk_size):
|
||
yield chunk
|
||
finally:
|
||
handle.close()
|
||
|
||
|
||
@router.post("/export/images")
|
||
async def export_images(req: ExportImagesRequest):
|
||
settings = get_settings()
|
||
if len(req.images) > settings.export_max_images:
|
||
raise HTTPException(status_code=413, detail=f"一次最多导出 {settings.export_max_images} 张图片，请分批导出")
|
||
|
||
root = _clean(req.title) or "采集图片"
|
||
# 小 ZIP 留在内存，大 ZIP 自动滚动到临时磁盘；响应阶段按块传输。
|
||
archive = tempfile.SpooledTemporaryFile(
|
||
max_size=settings.export_zip_memory_threshold_bytes,
|
||
mode="w+b",
|
||
)
|
||
used: set[str] = set()
|
||
total_source_bytes = 0
|
||
|
||
try:
|
||
with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as zf:
|
||
for img in req.images:
|
||
try:
|
||
data, ctype = await _download(img.url)
|
||
except UnsafeRemoteUrlError as exc:
|
||
raise HTTPException(status_code=400, detail=f"图片地址不允许导出：{exc}") from exc
|
||
except ResourceSizeLimitError as exc:
|
||
raise HTTPException(status_code=413, detail=f"无法导出「{img.key or '图片'}」：{exc}") from exc
|
||
except ResourceDownloadTimeoutError as exc:
|
||
raise HTTPException(status_code=504, detail=str(exc)) from exc
|
||
except Exception: # noqa: BLE001
|
||
continue # 单张下载失败不中断整包
|
||
if not _is_image(img.url, ctype):
|
||
continue
|
||
total_source_bytes += len(data)
|
||
if total_source_bytes > settings.export_zip_max_bytes:
|
||
raise HTTPException(
|
||
status_code=413,
|
||
detail=f"导出图片总大小超过 {settings.export_zip_max_bytes // 1024 // 1024} MB，请减少勾选图片后重试",
|
||
)
|
||
|
||
ext = _ext(img.url, ctype)
|
||
base = _clean(img.key) or "image"
|
||
if img.variantName:
|
||
base += f"-{_clean(img.variantName)}"
|
||
filename = f"{base}{ext}"
|
||
if filename in used: # 同名加序号防覆盖
|
||
stem = filename[: -len(ext)]
|
||
n = 2
|
||
while f"{stem}-{n}{ext}" in used:
|
||
n += 1
|
||
filename = f"{stem}-{n}{ext}"
|
||
used.add(filename)
|
||
|
||
group = _clean(img.groupName) or "图片"
|
||
zf.writestr(f"{root}/{group}/{filename}", data)
|
||
if archive.tell() > settings.export_zip_max_bytes:
|
||
raise HTTPException(
|
||
status_code=413,
|
||
detail=f"导出 ZIP 超过 {settings.export_zip_max_bytes // 1024 // 1024} MB，请减少勾选图片后重试",
|
||
)
|
||
except Exception:
|
||
archive.close()
|
||
raise
|
||
|
||
archive.seek(0)
|
||
return StreamingResponse(
|
||
_stream_file(archive),
|
||
media_type="application/zip",
|
||
headers={"Content-Disposition": 'attachment; filename="collect.zip"'},
|
||
)
|
||
|
||
|
||
@router.post("/export/image")
|
||
async def export_image(img: ExportImageItem):
|
||
"""代理下载单张图片，兼容源站防盗链与本地媒体路径。"""
|
||
try:
|
||
data, ctype = await _download(img.url)
|
||
except UnsafeRemoteUrlError as exc:
|
||
raise HTTPException(status_code=400, detail=f"图片地址不允许下载：{exc}") from exc
|
||
except ResourceSizeLimitError as exc:
|
||
raise HTTPException(status_code=413, detail=f"图片文件过大，无法下载：{exc}") from exc
|
||
except ResourceDownloadTimeoutError as exc:
|
||
raise HTTPException(status_code=504, detail=str(exc)) from exc
|
||
except Exception as exc: # noqa: BLE001
|
||
raise HTTPException(status_code=422, detail="图片下载失败") from exc
|
||
if not _is_image(img.url, ctype):
|
||
raise HTTPException(status_code=400, detail="该地址不是图片")
|
||
|
||
ext = _ext(img.url, ctype)
|
||
base = _clean(img.key) or "image"
|
||
if img.variantName:
|
||
base += f"-{_clean(img.variantName)}"
|
||
return Response(
|
||
content=data,
|
||
media_type=ctype,
|
||
headers={"Content-Disposition": _download_disposition(f"{base}{ext}")},
|
||
)
|