94 lines
3.0 KiB
Python
94 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import imghdr
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
from tempfile import NamedTemporaryFile
|
|
from time import time
|
|
|
|
from fastapi import UploadFile
|
|
from PIL import Image, ImageFilter
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.user import User
|
|
|
|
|
|
ALLOWED_IMAGE_TYPES = {"jpeg", "png"}
|
|
UPLOAD_ROOT = Path(__file__).resolve().parents[2] / "static" / "uploads"
|
|
|
|
|
|
def _validate_image(content: bytes) -> str:
|
|
image_type = imghdr.what(None, h=content)
|
|
if image_type not in ALLOWED_IMAGE_TYPES:
|
|
raise ValueError("仅支持 jpg/png 图片")
|
|
return image_type
|
|
|
|
|
|
def _save_optimized_variants(content: bytes) -> tuple[str, str, str]:
|
|
_validate_image(content)
|
|
|
|
with Image.open(BytesIO(content)) as image:
|
|
image = image.convert("RGB")
|
|
image.thumbnail((800, 800))
|
|
|
|
with NamedTemporaryFile(delete=False, suffix=".jpg") as main_file:
|
|
image.save(main_file.name, format="JPEG", quality=85)
|
|
main_path = main_file.name
|
|
|
|
blur_image = image.filter(ImageFilter.GaussianBlur(radius=15))
|
|
with NamedTemporaryFile(delete=False, suffix="_blur.jpg") as blur_file:
|
|
blur_image.save(blur_file.name, format="JPEG", quality=80)
|
|
blur_path = blur_file.name
|
|
|
|
timestamp = str(int(time()))
|
|
return main_path, blur_path, timestamp
|
|
|
|
|
|
def _persist_avatar_files(user_id: int, timestamp: str, main_path: str, blur_path: str) -> tuple[str, str]:
|
|
user_dir = UPLOAD_ROOT / "avatars" / str(user_id)
|
|
user_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
avatar_name = f"{timestamp}.jpg"
|
|
blur_name = f"{timestamp}_blur.jpg"
|
|
avatar_path = user_dir / avatar_name
|
|
blur_target_path = user_dir / blur_name
|
|
|
|
avatar_path.write_bytes(Path(main_path).read_bytes())
|
|
blur_target_path.write_bytes(Path(blur_path).read_bytes())
|
|
|
|
return avatar_name, blur_name
|
|
|
|
|
|
async def upload_user_avatar(session: AsyncSession, user: User, file: UploadFile) -> dict:
|
|
content = await file.read()
|
|
if not content:
|
|
raise ValueError("上传文件不能为空")
|
|
if len(content) > 5 * 1024 * 1024:
|
|
raise ValueError("图片大小不能超过 5MB")
|
|
|
|
main_path, blur_path, timestamp = _save_optimized_variants(content)
|
|
|
|
try:
|
|
avatar_name, blur_name = _persist_avatar_files(user.id, timestamp, main_path, blur_path)
|
|
avatar_url = f"/static/uploads/avatars/{user.id}/{avatar_name}"
|
|
blur_url = f"/static/uploads/avatars/{user.id}/{blur_name}"
|
|
|
|
user.avatar_url = avatar_url
|
|
user.avatar_blur_url = blur_url
|
|
|
|
session.add(user)
|
|
await session.commit()
|
|
await session.refresh(user)
|
|
|
|
return {
|
|
"avatar_url": user.avatar_url,
|
|
"avatar_blur_url": user.avatar_blur_url,
|
|
"local_files": {
|
|
"avatar": avatar_name,
|
|
"blur": blur_name,
|
|
},
|
|
}
|
|
finally:
|
|
Path(main_path).unlink(missing_ok=True)
|
|
Path(blur_path).unlink(missing_ok=True)
|