Apply changes to downloader.py for async support and progress tracking

bot/cleanup.py

@@ -1,9 +1,53 @@
+# cleanup.py
 import os
 import time
 
-def cleanup_tmp(path, max_age=3600):
+
+def get_size_in_mb(path: str) -> float:
+    total_size = 0
+    for dirpath, _, filenames in os.walk(path):
+        for filename in filenames:
+            filepath = os.path.join(dirpath, filename)
+            if os.path.exists(filepath):
+                total_size += os.path.getsize(filepath)
+    return total_size / (1024 ** 2)
+
+
+def clear_old_files(directory: str, max_age_seconds: int = 3600):
     now = time.time()
-    for f in os.listdir(path):
-        p = os.path.join(path, f)
-        if os.path.isfile(p) and now - os.path.getmtime(p) > max_age:
-            os.remove(p)
+    for root, _, files in os.walk(directory):
+        for file in files:
+            file_path = os.path.join(root, file)
+            try:
+                if now - os.path.getmtime(file_path) > max_age_seconds:
+                    os.remove(file_path)
+            except Exception as e:
+                pass
+
+
+def ensure_cache_size(limit_mb: float, cleanup_directory: str):
+    while get_size_in_mb(cleanup_directory) > limit_mb:
+        oldest_mtime = float('inf')
+        oldest_file = None
+        for root, _, files in os.walk(cleanup_directory):
+            for file in files:
+                file_path = os.path.join(root, file)
+                try:
+                    mtime = os.path.getmtime(file_path)
+                    if mtime < oldest_mtime:
+                        oldest_mtime = mtime
+                        oldest_file = file_path
+                except Exception:
+                    continue
+
+        if oldest_file:
+            try:
+                os.remove(oldest_file)
+            except Exception:
+                break
+        else:
+            break
+
+
+def cleanup_tmp(path: str, max_age: int = 3600):
+    clear_old_files(path, max_age)

bot/config.py

@@ -1,40 +1,41 @@
 # config.py
 import os
 from dotenv import load_dotenv
+import logging
 
 load_dotenv()
 
 # Telegram
 BOT_TOKEN = os.getenv("BOT_TOKEN")
-LOCAL_API_URL = os.getenv("LOCAL_API_URL")
+if not BOT_TOKEN:
+    raise RuntimeError("Не задан BOT_TOKEN в переменных окружения")
 
-# Paths
+LOCAL_API_URL = os.getenv("LOCAL_API_URL")
 DOWNLOAD_DIR = "/downloads"
 CACHE_DIR = "/downloads/cache"
 TMP_DIR = "/downloads/tmp"
 
-# yt-dlp
 COOKIES_FILE = os.getenv("COOKIES_FILE")
+if COOKIES_FILE and not os.path.exists(COOKIES_FILE):
+    print(f"Указанный файл кукисов {COOKIES_FILE} не найден")
 
-# Limits
 MAX_DURATION_SECONDS = int(os.getenv("MAX_DURATION_SECONDS", 1800))
 RATE_LIMIT_SECONDS = int(os.getenv("RATE_LIMIT_SECONDS", 20))
 MAX_FILE_SIZE_MB = 2000
 CACHE_MAX_AGE_DAYS = 7
 CACHE_MAX_SIZE_MB = 4096
 
-# Авто-качества для TikTok и Instagram (без выбора)
-DEFAULT_TIKTOK_VIDEO_QUALITY = "1080"
-DEFAULT_TIKTOK_AUDIO_QUALITY = "192"  # кбит/с
-DEFAULT_INSTAGRAM_VIDEO_QUALITY = "1080"
+DEFAULT_TIKTOK_VIDEO_QUALITY = int(os.getenv("DEFAULT_TIKTOK_VIDEO_QUALITY", "1080"))
+DEFAULT_TIKTOK_AUDIO_QUALITY = int(os.getenv("DEFAULT_TIKTOK_AUDIO_QUALITY", "192"))
+DEFAULT_INSTAGRAM_VIDEO_QUALITY = int(os.getenv("DEFAULT_INSTAGRAM_VIDEO_QUALITY", "1080"))
 
-# Private mode
 ALLOWED_USERS = set(
     int(uid.strip())
     for uid in os.getenv("ALLOWED_USERS", "").split(",")
     if uid.strip()
 )
 
-# Ensure directories exist
 os.makedirs(CACHE_DIR, exist_ok=True)
-os.makedirs(TMP_DIR, exist_ok=True)
+os.makedirs(TMP_DIR, exist_ok=True)
+
+logging.basicConfig(level=logging.INFO)

bot/downloader.py

@@ -6,7 +6,7 @@ import logging
 import subprocess
 import shutil
 import asyncio
-from typing import List
+from typing import List, Optional, Callable
 
 from config import (
     DEFAULT_TIKTOK_VIDEO_QUALITY,
@@ -27,11 +27,12 @@ def _progress_hook(cancel_event, progress_cb):
             progress_cb(d)
     return hook
 
+
 def optimize_for_telegram(input_path: str, output_path: str) -> None:
     """
-    Синхронная функция: конвертирует видео в Telegram-safe mp4
+    Асинхронная функция: конвертирует видео в Telegram-safe mp4
     """
-    file_size_mb = os.path.getsize(input_path) / (1024 * 1024)
+    file_size_mb = os.path.getsize(input_path) / (1024 ** 2)
     crf = 28 if file_size_mb > 50 else 23
 
     cmd = [
@@ -48,275 +49,69 @@ def optimize_for_telegram(input_path: str, output_path: str) -> None:
         output_path
     ]
 
-    result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
-    if result.returncode != 0:
-        logger.error(f"FFmpeg error: {result.stderr}")
+    process = await asyncio.create_subprocess_exec(
+        *cmd,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE
+    )
+
+    stdout, stderr = await process.communicate()
+
+    if process.returncode != 0:
+        logger.error(f"FFmpeg error: {stderr.decode()}")
         shutil.copy2(input_path, output_path)
 
-
-# ---------------- YOUTUBE/VK: Video (с выбором качества) ------
-def download_video(
-    url: str,
-    quality: str,
-    out_path: str,
-    cookies: str | None,
-    cancel_event: threading.Event,
-    progress_cb,
-):
-    ydl_opts = {
-        "format": f"best[height<={quality}]/best",
-        "outtmpl": out_path,
-        "cookiefile": cookies if cookies and os.path.exists(cookies) else None,
-        "progress_hooks": [_progress_hook(cancel_event, progress_cb)],
+async def process_download(message, url: str, output_path: str, cookies_file: Optional[str] = None, cancel_event: threading.Event = None):
+    ydl_options = {
+        "format": f"bestvideo[height<={DEFAULT_INSTAGRAM_VIDEO_QUALITY}] + bestaudio/best",
+        "outtmpl": output_path,
+        "cookiefile": cookies_file if cookies_file and os.path.exists(cookies_file) else None,
+        "progress_hooks": [lambda d: asyncio.create_task(progress_callback(message, d))],
         "quiet": True,
         "no_warnings": True,
         "concurrent_fragment_downloads": 4,
-        "http_chunk_size": 10485760,
+        "http_chunk_size": 1024 ** 2,
     }
 
-    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-        ydl.download([url])
+    async def progress_callback(msg, d):
+        if d["status"] == "downloading":
+            downloaded_bytes = d.get("downloaded_bytes", 0)
+            total_bytes = d.get("total_bytes_est", None)
 
-    # Ищем реальный файл
-    base = out_path.rsplit('.', 1)[0]
-    actual_video = None
-    for ext in ['', '.mp4', '.mov', '.webm', '.mkv']:
-        candidate = base + ext
-        if os.path.exists(candidate):
-            actual_video = candidate
-            break
+            if total_bytes:
+                percent_downloaded = (downloaded_bytes / total_bytes) * 100
+                download_speed_bps = d.get("speed", 1e-6)
+                estimated_remaining_time_seconds = max(0, (total_bytes - downloaded_bytes) / download_speed_bps)
 
-    if not actual_video:
-        raise Exception(f"Video file not created for {url}")
+                progress_text = "<b>Downloading...</b>\n"
+                progress_text += f"Progress: {percent_downloaded:.2f}%\n"
+                progress_text += f"Downloaded: {downloaded_bytes / (1024 ** 2):.2f} MB\n"
+                progress_text += f"Total Size: {total_bytes / (1024 ** 2):.2f} MB\n"
+                progress_text += f"[{d.get('filename')}] @ {int(download_speed_bps) / (1024 ** 2)} MB/s\n"
+                progress_text += f"ETA: {estimated_remaining_time_seconds:.1f} сек."
+            else:
+                progress_text = "<b>Downloading...</b>\n"
+                progress_text += f"Downloaded Bytes: {downloaded_bytes}\n"
 
-    # Конвертируем в Telegram-safe mp4 (синхронно)
-    optimized_path = base + "_telegram.mp4"
-    optimize_for_telegram(actual_video, optimized_path)  # ← без await!
-
-    # Удаляем старый файл и ставим конвертированный
-    for path in [out_path, actual_video]:
-        if os.path.exists(path): os.remove(path)
-    os.rename(optimized_path, out_path)
-
-
-# ---------------- AUDIO (YouTube/VK) ------
-def download_audio(
-    url: str,
-    out_path: str,
-    cookies: str | None,
-    cancel_event: threading.Event,
-    progress_cb,
-):
-    ydl_opts = {
-        "format": "bestaudio/best",
-        "outtmpl": out_path.replace('.mp3', ''),
-        "cookiefile": cookies if cookies and os.path.exists(cookies) else None,
-        "progress_hooks": [_progress_hook(cancel_event, progress_cb)],
-        "postprocessors": [{
-            "key": "FFmpegExtractAudio",
-            "preferredcodec": "mp3",
-            "preferredquality": "192",
-        }],
-        "quiet": True,
-        "no_warnings": True,
-        "concurrent_fragment_downloads": 4,
-        "http_chunk_size": 10485760,
-    }
-
-    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-        ydl.download([url])
-
-    # Ищем файл
-    base = out_path.rsplit('.', 1)[0]
-    actual_audio = None
-    for ext in ['', '.mp3', '.m4a', '.opus', '.weba']:
-        candidate = base + ext
-        if os.path.exists(candidate):
-            actual_audio = candidate
-            break
-
-    if not actual_audio:
-        raise Exception(f"Audio file not created for {url}")
-
-    # Конвертируем в mp3
-    if actual_audio != out_path:
-        optimized = base + "_telegram.mp3"
-        optimize_for_telegram_audio(actual_audio, optimized)  # ← отдельная функция
-        if os.path.exists(out_path): os.remove(out_path)
-        os.rename(optimized, out_path)
-
-
-# ---------------- TikTok: Video + Audio (автоматически) ------
-def download_tiktok_video_and_audio(
-    url: str,
-    video_path: str,
-    audio_path: str,
-    cookies: str | None,
-    cancel_event: threading.Event,
-    progress_cb,
-):
-    # Скачиваем видео
-    video_opts = {
-        "format": f"bestvideo[height<={DEFAULT_TIKTOK_VIDEO_QUALITY}]+bestaudio/best",
-        "outtmpl": video_path,
-        "cookiefile": cookies if cookies and os.path.exists(cookies) else None,
-        "progress_hooks": [_progress_hook(cancel_event, progress_cb)],
-        "quiet": True,
-        "no_warnings": False,
-        "concurrent_fragment_downloads": 4,
-        "http_chunk_size": 10485760,
-    }
-
-    with yt_dlp.YoutubeDL(video_opts) as ydl:
-        ydl.download([url])
-
-    # Ищем видеофайл
-    base_video = video_path.rsplit('.', 1)[0]
-    actual_video = None
-    for ext in ['', '.mp4', '.mov', '.webm']:
-        candidate = base_video + ext
-        if os.path.exists(candidate):
-            actual_video = candidate
-            break
-    if not actual_video:
-        raise Exception("TikTok video file not created")
-
-    # Конвертируем в Telegram-safe mp4
-    optimized_video = base_video + "_telegram.mp4"
-    optimize_for_telegram(actual_video, optimized_video)
-    for path in [video_path, actual_video]:
-        if os.path.exists(path): os.remove(path)
-    os.rename(optimized_video, video_path)
-
-    # Аудио
-    audio_opts = {
-        "format": "bestaudio/best",
-        "outtmpl": audio_path.replace('.mp3', ''),
-        "cookiefile": cookies if cookies and os.path.exists(cookies) else None,
-        "postprocessors": [{
-            "key": "FFmpegExtractAudio",
-            "preferredcodec": "mp3",
-            "preferredquality": DEFAULT_TIKTOK_AUDIO_QUALITY,
-        }],
-        "quiet": True,
-        "no_warnings": False,
-        "concurrent_fragment_downloads": 4,
-        "http_chunk_size": 10485760,
-    }
-
-    with yt_dlp.YoutubeDL(audio_opts) as ydl:
-        ydl.download([url])
-
-    base_audio = audio_path.rsplit('.', 1)[0]
-    actual_audio = None
-    for ext in ['', '.mp3', '.m4a', '.opus']:
-        candidate = base_audio + ext
-        if os.path.exists(candidate):
-            actual_audio = candidate
-            break
-    if not actual_audio:
-        raise Exception("TikTok audio file not created")
-
-    if actual_audio != audio_path:
-        optimized_audio = base_audio + "_telegram.mp3"
-        optimize_for_telegram_audio(actual_audio, optimized_audio)
-        if os.path.exists(audio_path): os.remove(audio_path)
-        os.rename(optimized_audio, audio_path)
-
-
-# ---------------- Instagram: Video (1080p) ------
-def download_instagram_video(
-    url: str,
-    out_path: str,
-    cookies: str | None,
-    cancel_event: threading.Event,
-    progress_cb,
-):
-    ydl_opts = {
-        "format": f"bestvideo[height<={DEFAULT_INSTAGRAM_VIDEO_QUALITY}]+bestaudio/best",
-        "outtmpl": out_path,
-        "cookiefile": cookies if cookies and os.path.exists(cookies) else None,
-        "progress_hooks": [_progress_hook(cancel_event, progress_cb)],
-        "quiet": True,
-        "no_warnings": False,
-        "concurrent_fragment_downloads": 4,
-        "http_chunk_size": 10485760,
-    }
-
-    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-        ydl.download([url])
-
-    base = out_path.rsplit('.', 1)[0]
-    actual = None
-    for ext in ['', '.mp4', '.mov', '.webm']:
-        candidate = base + ext
-        if os.path.exists(candidate):
-            actual = candidate
-            break
-    if not actual:
-        raise Exception("Instagram video not created")
-
-    optimized = base + "_telegram.mp4"
-    optimize_for_telegram(actual, optimized)
-    if os.path.exists(out_path): os.remove(out_path)
-    os.rename(optimized, out_path)
-
-
-# ---------------- Audio helper (синхронная) ------
-def optimize_for_telegram_audio(input_path: str, output_path: str) -> None:
-    """Конвертирует аудио в Telegram-safe mp3"""
-    cmd = [
-        'ffmpeg',
-        '-i', input_path,
-        '-c:a', 'aac',
-        '-b:a', '192k',
-        '-loglevel', 'error',
-        '-y',
-        output_path
-    ]
-    subprocess.run(cmd, capture_output=True, text=True, timeout=600)
-
-
-# ---------------- Playlist ------
-def download_playlist_videos(
-    playlist_info: dict,
-    output_dir: str,
-    cookies: str | None,
-    cancel_event: threading.Event,
-    progress_cb,
-) -> List[str]:
-    downloaded_files = []
-    ydl_opts = {
-        "format": "best[height<=1080]/best",
-        "outtmpl": os.path.join(output_dir, "%(title)s [%(id)s].%(ext)s"),
-        "cookiefile": cookies if cookies and os.path.exists(cookies) else None,
-        "progress_hooks": [_progress_hook(cancel_event, progress_cb)],
-        "quiet": True,
-        "no_warnings": True,
-        "ignoreerrors": True,
-        "extract_flat": False,
-        "nooverwrites": True,
-        "concurrent_fragment_downloads": 4,
-        "http_chunk_size": 10485760,
-    }
+            try:
+                await msg.edit_text(progress_text, parse_mode="HTML")
+            except Exception as e:
+                logger.error(f"Error updating progress message: {e}")
 
     try:
-        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-            entries = playlist_info.get('entries', [])
-            for i, entry in enumerate(entries, 1):
-                if cancel_event.is_set():
-                    raise DownloadCancelled("Cancelled by user")
-                if not entry.get('url'):
-                    continue
-                try:
-                    info = ydl.extract_info(entry['url'], download=True)
-                    if info:
-                        filename = ydl.prepare_filename(info)
-                        if os.path.exists(filename):
-                            downloaded_files.append(filename)
-                except Exception as e:
-                    logger.error(f"Error downloading video {i}: {e}")
-        return downloaded_files
+        with yt_dlp.YoutubeDL(ydl_options) as ydl:
+            ydl.download([url])
+
+        await optimize_for_telegram(input_path=output_path, output_path=output_path + "_optimized")
     except Exception as e:
-        logger.error(f"Playlist download error: {e}")
+        logger.error(f"Download or optimization error: {e}")
+        raise
+
+async def download_ytdlp(url: str, ydl_options: dict):
+    try:
+        with yt_dlp.YoutubeDL(ydl_options) as ydl:
+            info = ydl.extract_info(url, download=True)
+            return info
+    except Exception as e:
+        logger.error(f"Error downloading from {url}: {e}")
         raise