77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import hashlib
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Tuple
|
||
from urllib.parse import urlparse
|
||
|
||
import httpx
|
||
|
||
# Корень хранилища картинок внутри контейнера/процесса
|
||
MEDIA_ROOT = Path(os.getenv("NYANIMEDB_MEDIA_ROOT", "media")).resolve()
|
||
|
||
|
||
def _guess_ext_from_url(url: str) -> str:
|
||
path = urlparse(url).path
|
||
_, ext = os.path.splitext(path)
|
||
if ext and len(ext) <= 5:
|
||
return ext
|
||
return ".jpg"
|
||
|
||
|
||
def _build_rel_path_from_hash(h: str, ext: str, subdir: str = "posters") -> Tuple[str, Path]:
|
||
"""
|
||
Строим путь вида subdir/ab/cd/<hash>.ext по sha1-хешу содержимого.
|
||
"""
|
||
level1 = h[:2]
|
||
level2 = h[2:4]
|
||
rel = f"{subdir}/{level1}/{level2}/{h}{ext}"
|
||
abs_path = MEDIA_ROOT / rel
|
||
return rel, abs_path
|
||
|
||
|
||
async def _fetch_bytes(url: str) -> bytes:
|
||
async with httpx.AsyncClient(timeout=20.0) as client:
|
||
r = await client.get(url)
|
||
r.raise_for_status()
|
||
return r.content
|
||
|
||
|
||
async def ensure_image_downloaded(url: str, subdir: str = "posters") -> str:
|
||
"""
|
||
Гарантирует, что картинка по URL лежит в MEDIA_ROOT/subdir в структуре:
|
||
subdir/ab/cd/<sha1(content)>.ext
|
||
|
||
Возвращает относительный путь (для записи в БД).
|
||
Один и тот же файл (по содержимому) всегда даёт один и тот же путь,
|
||
даже если URL меняется.
|
||
"""
|
||
# Скачиваем данные
|
||
data = await _fetch_bytes(url)
|
||
|
||
# Хешируем именно содержимое, а не URL
|
||
h = hashlib.sha1(data).hexdigest()
|
||
ext = _guess_ext_from_url(url)
|
||
|
||
rel, abs_path = _build_rel_path_from_hash(h, ext, subdir=subdir)
|
||
|
||
# Если файл уже есть (другой процесс/воркер успел сохранить) — просто возвращаем путь
|
||
if abs_path.exists():
|
||
return rel
|
||
|
||
abs_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Пишем во временный файл и затем делаем atomic rename
|
||
tmp_path = abs_path.with_suffix(abs_path.suffix + ".tmp")
|
||
|
||
def _write() -> None:
|
||
with open(tmp_path, "wb") as f:
|
||
f.write(data)
|
||
# os.replace атомарно заменит файл, даже если он уже появился
|
||
os.replace(tmp_path, abs_path)
|
||
|
||
await asyncio.to_thread(_write)
|
||
|
||
return rel
|