402 lines
16 KiB
Python
402 lines
16 KiB
Python
"""Логика скачивания.
|
||
|
||
Две стратегии за общим интерфейсом:
|
||
- HttpxDownloader — прямые ссылки на файлы, прогресс по chunk'ам.
|
||
- YtDlpDownloader — видео/медиа с сайтов (YouTube и сотни других) через yt-dlp.
|
||
|
||
`pick_downloader()` выбирает стратегию по URL. Прогресс отдаётся через
|
||
async-колбэк on_progress, который дёргается не чаще, чем раз в progress_interval.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import functools
|
||
import os
|
||
import re
|
||
import time
|
||
from collections.abc import Awaitable, Callable
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import BinaryIO
|
||
from urllib.parse import unquote, urlparse
|
||
|
||
import httpx
|
||
|
||
from app.config import settings
|
||
|
||
# Расширения, которые качаем напрямую через httpx, а не через yt-dlp.
|
||
_DIRECT_EXTENSIONS = {
|
||
".zip", ".rar", ".7z", ".tar", ".gz", ".tgz", ".bz2", ".xz",
|
||
".iso", ".dmg", ".exe", ".msi", ".apk", ".deb", ".rpm", ".appimage",
|
||
".pdf", ".epub", ".mobi", ".djvu",
|
||
".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg", ".bmp", ".tiff",
|
||
".mp3", ".flac", ".wav", ".ogg", ".m4a", ".aac",
|
||
".mp4", ".mkv", ".webm", ".mov", ".avi", ".flv", ".m4v", # прямые ссылки на медиа
|
||
".doc", ".docx", ".xls", ".xlsx", ".ppt", ".pptx", ".csv", ".txt", ".json",
|
||
".bin", ".img", ".dll", ".so",
|
||
}
|
||
|
||
OnProgress = Callable[["Progress"], Awaitable[None]]
|
||
|
||
|
||
@dataclass
|
||
class Progress:
|
||
downloaded_bytes: int
|
||
total_bytes: int | None
|
||
speed: float | None
|
||
eta: float | None
|
||
filename: str | None = None
|
||
|
||
|
||
@dataclass
|
||
class DownloadResult:
|
||
filename: str
|
||
size_bytes: int | None
|
||
path: Path
|
||
|
||
|
||
@dataclass
|
||
class Resolved:
|
||
"""Что вернул резолвер сайта: куда реально идти за файлом.
|
||
|
||
headers — если для прямой ссылки нужен Referer/Cookie/авторизация.
|
||
filename — если сайт знает «правильное» имя (иначе возьмём из ответа).
|
||
"""
|
||
|
||
download_url: str
|
||
filename: str | None = None
|
||
headers: dict[str, str] | None = None
|
||
|
||
|
||
def _safe_filename(name: str) -> str:
|
||
"""Чистим имя файла от разделителей пути и опасных символов."""
|
||
name = unquote(name).strip().replace("\\", "/").split("/")[-1]
|
||
name = re.sub(r'[<>:"/\\|?*\x00-\x1f]', "_", name)
|
||
name = name.strip(". ") or "download"
|
||
return name[:200]
|
||
|
||
|
||
def _unique_path(directory: Path, filename: str) -> Path:
|
||
"""Свободное имя по-браузерному: file.ext, file (1).ext, file (2).ext, ...
|
||
|
||
Best-effort: только подбирает имя, не резервирует его. Для случаев, где файл
|
||
пишет внешний инструмент (yt-dlp) и держать дескриптор нельзя.
|
||
"""
|
||
candidate = directory / filename
|
||
if not candidate.exists():
|
||
return candidate
|
||
stem, suffix = candidate.stem, candidate.suffix
|
||
i = 1
|
||
while True:
|
||
candidate = directory / f"{stem} ({i}){suffix}"
|
||
if not candidate.exists():
|
||
return candidate
|
||
i += 1
|
||
|
||
|
||
def _create_unique(directory: Path, filename: str) -> tuple[Path, BinaryIO]:
|
||
"""Атомарно создать НОВЫЙ файл, разводя дубли как браузер: name (1).ext и т.д.
|
||
|
||
Эксклюзивное создание (режим "xb") закрывает гонку между воркерами: имя
|
||
не просто подобрано, а сразу занято — параллельная загрузка не затрёт.
|
||
"""
|
||
directory.mkdir(parents=True, exist_ok=True)
|
||
candidate = directory / filename
|
||
stem, suffix = candidate.stem, candidate.suffix
|
||
i = 0
|
||
while True:
|
||
try:
|
||
return candidate, candidate.open("xb")
|
||
except FileExistsError:
|
||
i += 1
|
||
candidate = directory / f"{stem} ({i}){suffix}"
|
||
|
||
|
||
def is_direct_file(url: str) -> bool:
|
||
path = urlparse(url).path.lower()
|
||
return any(path.endswith(ext) for ext in _DIRECT_EXTENSIONS)
|
||
|
||
|
||
def _filename_from_response(resp: httpx.Response) -> str:
|
||
cd = resp.headers.get("content-disposition", "")
|
||
match = re.search(r"filename\*=(?:UTF-8'')?([^;]+)|filename=\"?([^\";]+)\"?", cd)
|
||
if match:
|
||
raw = match.group(1) or match.group(2)
|
||
if raw:
|
||
return _safe_filename(raw)
|
||
name = urlparse(str(resp.url)).path
|
||
return _safe_filename(name) if name and name != "/" else "download"
|
||
|
||
|
||
async def _stream_to_file(
|
||
url: str,
|
||
on_progress: OnProgress,
|
||
*,
|
||
headers: dict[str, str] | None = None,
|
||
filename: str | None = None,
|
||
) -> DownloadResult:
|
||
"""Общее ядро скачивания по прямой ссылке: стрим по chunk'ам, прогресс,
|
||
(n)-имена. Используется и httpx-загрузчиком, и резолверами сайтов."""
|
||
directory = settings.download_dir
|
||
directory.mkdir(parents=True, exist_ok=True)
|
||
timeout = httpx.Timeout(30.0, read=None)
|
||
async with httpx.AsyncClient(
|
||
follow_redirects=True, timeout=timeout, headers=headers
|
||
) as client:
|
||
async with client.stream("GET", url) as resp:
|
||
resp.raise_for_status()
|
||
name = _safe_filename(filename) if filename else _filename_from_response(resp)
|
||
total = (
|
||
int(resp.headers["content-length"])
|
||
if "content-length" in resp.headers
|
||
else None
|
||
)
|
||
target, fh = _create_unique(directory, name)
|
||
|
||
downloaded = 0
|
||
start = time.monotonic()
|
||
last_emit = 0.0
|
||
try:
|
||
with fh:
|
||
async for chunk in resp.aiter_bytes(settings.chunk_size):
|
||
fh.write(chunk)
|
||
downloaded += len(chunk)
|
||
now = time.monotonic()
|
||
if now - last_emit >= settings.progress_interval:
|
||
elapsed = now - start
|
||
speed = downloaded / elapsed if elapsed > 0 else None
|
||
eta = (
|
||
(total - downloaded) / speed
|
||
if total and speed and speed > 0
|
||
else None
|
||
)
|
||
await on_progress(
|
||
Progress(downloaded, total, speed, eta, target.name)
|
||
)
|
||
last_emit = now
|
||
except BaseException:
|
||
target.unlink(missing_ok=True)
|
||
raise
|
||
|
||
elapsed = time.monotonic() - start
|
||
speed = downloaded / elapsed if elapsed > 0 else None
|
||
await on_progress(Progress(downloaded, total or downloaded, speed, 0, target.name))
|
||
return DownloadResult(target.name, downloaded, target)
|
||
|
||
|
||
class UnsupportedURLError(Exception):
|
||
"""Ни одна стратегия не берётся за URL: не прямой файл и не из загрузчиков."""
|
||
|
||
|
||
# --- Реестр стратегий ---------------------------------------------------------
|
||
# Расширения регистрируются декоратором @register. pick_downloader() выбирает
|
||
# первую подходящую по убыванию priority. Кастомные экстракторы (priority > 0)
|
||
# перехватывают URL раньше встроенных httpx/yt-dlp.
|
||
|
||
_REGISTRY: list[type["Downloader"]] = []
|
||
|
||
|
||
def register(cls: type["Downloader"]) -> type["Downloader"]:
|
||
_REGISTRY.append(cls)
|
||
return cls
|
||
|
||
|
||
class Downloader:
|
||
"""Базовый интерфейс стратегии. priority — кто раньше перехватывает URL
|
||
(больше = раньше); matches() — берётся ли эта стратегия за данный URL.
|
||
|
||
label/kind — для страницы «Загрузчики» (kind: direct | extractor | fallback)."""
|
||
|
||
priority: int = 0
|
||
label: str = "" # человекочитаемое имя; пусто -> берём имя класса
|
||
kind: str = "other"
|
||
|
||
@classmethod
|
||
def matches(cls, url: str) -> bool:
|
||
return False
|
||
|
||
async def download(self, url: str, on_progress: OnProgress) -> DownloadResult:
|
||
raise NotImplementedError
|
||
|
||
|
||
class SiteExtractor(Downloader):
|
||
"""Удобная база для сайтов: реализуй matches() и resolve() — достать прямую
|
||
ссылку (и при нужде Referer/Cookie). Скачивание, прогресс и (n)-имена общие.
|
||
|
||
Если сайту нужен нестандартный процесс (HLS, сегменты) — наследуйся прямо от
|
||
Downloader и переопредели download() целиком."""
|
||
|
||
kind = "extractor"
|
||
|
||
async def resolve(self, url: str) -> Resolved:
|
||
raise NotImplementedError
|
||
|
||
async def download(self, url: str, on_progress: OnProgress) -> DownloadResult:
|
||
resolved = await self.resolve(url)
|
||
return await _stream_to_file(
|
||
resolved.download_url,
|
||
on_progress,
|
||
headers=resolved.headers,
|
||
filename=resolved.filename,
|
||
)
|
||
|
||
|
||
@register
|
||
class HttpxDownloader(Downloader):
|
||
"""Прямые ссылки на файлы (.zip, .pdf, .mp4 …) — стрим по chunk'ам."""
|
||
|
||
priority = 10
|
||
label = "Прямые файлы"
|
||
kind = "direct"
|
||
|
||
@classmethod
|
||
def matches(cls, url: str) -> bool:
|
||
return is_direct_file(url)
|
||
|
||
async def download(self, url: str, on_progress: OnProgress) -> DownloadResult:
|
||
return await _stream_to_file(url, on_progress)
|
||
|
||
|
||
@functools.lru_cache(maxsize=1)
|
||
def _ytdlp_site_extractors() -> tuple:
|
||
"""Профильные экстракторы yt-dlp без generic (он матчит почти любой http-URL).
|
||
|
||
Грузим один раз и кешируем — список большой и тянется лениво.
|
||
"""
|
||
from yt_dlp.extractor import gen_extractor_classes
|
||
|
||
return tuple(ie for ie in gen_extractor_classes() if ie.IE_NAME != "generic")
|
||
|
||
|
||
def _ytdlp_supports(url: str) -> bool:
|
||
"""Есть ли у yt-dlp профильный экстрактор под этот URL (generic не в счёт)."""
|
||
try:
|
||
return any(ie.suitable(url) for ie in _ytdlp_site_extractors())
|
||
except Exception: # noqa: BLE001 — проблемы yt-dlp не должны ронять выбор
|
||
return False
|
||
|
||
|
||
def _ytdlp_cookie_opts(from_browser: str | None, cookies_file: Path | str | None) -> dict:
|
||
"""Опции cookies для yt-dlp из настроек — для сайтов, блокирующих анонимов."""
|
||
opts: dict = {}
|
||
if cookies_file:
|
||
opts["cookiefile"] = str(cookies_file)
|
||
if from_browser:
|
||
browser, _, profile = from_browser.partition(":")
|
||
# yt-dlp ждёт кортеж (browser, profile, keyring, container).
|
||
opts["cookiesfrombrowser"] = (browser.strip(), profile.strip() or None, None, None)
|
||
return opts
|
||
|
||
|
||
@register
|
||
class YtDlpDownloader(Downloader):
|
||
"""yt-dlp: YouTube и сотни сайтов. Берётся за URL последним и только если у
|
||
yt-dlp есть профильный экстрактор под него (generic-угадывание не считаем —
|
||
иначе «ловит всё» и ошибки про неподдерживаемую ссылку не будет).
|
||
|
||
yt-dlp синхронный, поэтому крутим его в потоке, а progress-хуки прокидываем
|
||
обратно в event loop через run_coroutine_threadsafe."""
|
||
|
||
priority = -100
|
||
label = "yt-dlp (видео и сайты)"
|
||
kind = "fallback"
|
||
|
||
@classmethod
|
||
def matches(cls, url: str) -> bool:
|
||
return _ytdlp_supports(url)
|
||
|
||
def __init__(self) -> None:
|
||
self._dir = settings.download_dir
|
||
self._interval = settings.progress_interval
|
||
|
||
async def download(self, url: str, on_progress: OnProgress) -> DownloadResult:
|
||
self._dir.mkdir(parents=True, exist_ok=True)
|
||
loop = asyncio.get_running_loop()
|
||
last_emit = 0.0
|
||
result_holder: dict[str, object] = {}
|
||
|
||
def hook(d: dict) -> None:
|
||
nonlocal last_emit
|
||
status = d.get("status")
|
||
if status == "downloading":
|
||
now = time.monotonic()
|
||
if now - last_emit < self._interval:
|
||
return
|
||
last_emit = now
|
||
downloaded = d.get("downloaded_bytes") or 0
|
||
total = d.get("total_bytes") or d.get("total_bytes_estimate")
|
||
progress = Progress(
|
||
downloaded_bytes=downloaded,
|
||
total_bytes=total,
|
||
speed=d.get("speed"),
|
||
eta=d.get("eta"),
|
||
filename=os.path.basename(d.get("filename") or "") or None,
|
||
)
|
||
asyncio.run_coroutine_threadsafe(on_progress(progress), loop)
|
||
elif status == "finished":
|
||
result_holder["path"] = d.get("filename")
|
||
|
||
def run_blocking() -> DownloadResult:
|
||
# Импортируем лениво, чтобы httpx-only сценарий не тянул yt-dlp.
|
||
from yt_dlp import YoutubeDL
|
||
|
||
base_opts = {
|
||
"noprogress": True,
|
||
"quiet": True,
|
||
"no_warnings": True,
|
||
"noplaylist": True,
|
||
**_ytdlp_cookie_opts(settings.cookies_from_browser, settings.cookies_file),
|
||
}
|
||
default_tmpl = str(self._dir / "%(title)s [%(id)s].%(ext)s")
|
||
|
||
# Фаза 1: узнаём имя файла, не качая, чтобы развести дубли как браузер.
|
||
with YoutubeDL({**base_opts, "outtmpl": default_tmpl}) as probe:
|
||
info = probe.extract_info(url, download=False)
|
||
predicted = Path(probe.prepare_filename(info))
|
||
target = _unique_path(self._dir, predicted.name)
|
||
# %(ext)s оставляем yt-dlp (контейнер может смениться при склейке),
|
||
# а литеральную часть имени экранируем: % -> %% (вдруг в названии есть %).
|
||
stem = str(target.with_suffix("")).replace("%", "%%")
|
||
outtmpl = f"{stem}.%(ext)s"
|
||
|
||
# Фаза 2: качаем в выбранный путь.
|
||
with YoutubeDL({**base_opts, "outtmpl": outtmpl, "progress_hooks": [hook]}) as ydl:
|
||
info = ydl.extract_info(url, download=True)
|
||
final_path = result_holder.get("path") or ydl.prepare_filename(info)
|
||
path = Path(str(final_path))
|
||
size = path.stat().st_size if path.exists() else None
|
||
return DownloadResult(path.name, size, path)
|
||
|
||
return await loop.run_in_executor(None, run_blocking)
|
||
|
||
|
||
def pick_downloader(url: str) -> Downloader:
|
||
"""Первая подходящая стратегия по убыванию priority.
|
||
|
||
Кастомные экстракторы (priority > 0) перехватывают раньше httpx (10) и
|
||
yt-dlp (-100). Если не взялся никто — URL не поддерживается."""
|
||
for cls in sorted(_REGISTRY, key=lambda c: c.priority, reverse=True):
|
||
if cls.matches(url):
|
||
return cls()
|
||
raise UnsupportedURLError(
|
||
"Не могу скачать эту ссылку: это не прямой файл и ни один загрузчик "
|
||
"её не поддерживает."
|
||
)
|
||
|
||
|
||
def list_strategies() -> list[dict]:
|
||
"""Описание зарегистрированных стратегий для страницы «Загрузчики»."""
|
||
result: list[dict] = []
|
||
for cls in sorted(_REGISTRY, key=lambda c: c.priority, reverse=True):
|
||
doc = (cls.__doc__ or "").strip().split("\n")[0].strip()
|
||
result.append(
|
||
{
|
||
"name": cls.label or cls.__name__,
|
||
"kind": cls.kind,
|
||
"priority": cls.priority,
|
||
"description": doc or None,
|
||
}
|
||
)
|
||
return result
|