"""Логика скачивания. Две стратегии за общим интерфейсом: - HttpxDownloader — прямые ссылки на файлы, прогресс по chunk'ам. - YtDlpDownloader — видео/медиа с сайтов (YouTube и сотни других) через yt-dlp. `pick_downloader()` выбирает стратегию по URL. Прогресс отдаётся через async-колбэк on_progress, который дёргается не чаще, чем раз в progress_interval. """ from __future__ import annotations import asyncio import os import re import time from collections.abc import Awaitable, Callable from dataclasses import dataclass from pathlib import Path 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 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: """Не перезатираем существующие файлы — добавляем (1), (2), ...""" 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 is_direct_file(url: str) -> bool: path = urlparse(url).path.lower() return any(path.endswith(ext) for ext in _DIRECT_EXTENSIONS) class Downloader: """Базовый интерфейс стратегии скачивания.""" async def download(self, url: str, on_progress: OnProgress) -> DownloadResult: raise NotImplementedError class HttpxDownloader(Downloader): def __init__(self) -> None: self._dir = settings.download_dir self._chunk = settings.chunk_size self._interval = settings.progress_interval def _filename_from_response(self, url: str, 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 download(self, url: str, on_progress: OnProgress) -> DownloadResult: self._dir.mkdir(parents=True, exist_ok=True) timeout = httpx.Timeout(30.0, read=None) async with httpx.AsyncClient(follow_redirects=True, timeout=timeout) as client: async with client.stream("GET", url) as resp: resp.raise_for_status() filename = self._filename_from_response(url, resp) total = int(resp.headers["content-length"]) if "content-length" in resp.headers else None target = _unique_path(self._dir, filename) downloaded = 0 start = time.monotonic() last_emit = 0.0 try: with target.open("wb") as fh: async for chunk in resp.aiter_bytes(self._chunk): fh.write(chunk) downloaded += len(chunk) now = time.monotonic() if now - last_emit >= self._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 YtDlpDownloader(Downloader): """Скачивание через yt-dlp. yt-dlp синхронный, поэтому крутим его в потоке, а progress-хуки прокидываем обратно в event loop через call_soon_threadsafe.""" 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 ydl_opts = { "outtmpl": str(self._dir / "%(title)s [%(id)s].%(ext)s"), "progress_hooks": [hook], "noprogress": True, "quiet": True, "no_warnings": True, "noplaylist": True, } with YoutubeDL(ydl_opts) as ydl: info = ydl.extract_info(url, download=True) final_path = result_holder.get("path") if not final_path: final_path = 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: """Прямые ссылки на файлы — httpx, всё остальное (страницы сайтов) — yt-dlp.""" return HttpxDownloader() if is_direct_file(url) else YtDlpDownloader()