Update DragMusic/platforms/Youtube.py

DragMusic/platforms/Youtube.py

@@ -2,7 +2,7 @@ import asyncio
 import os
 import re
 import json
-import httpx  # Make sure to import httpx
+import httpx
 from typing import Union
 import traceback
 import yt_dlp
@@ -11,6 +11,9 @@ from pyrogram.types import Message
 from youtubesearchpython.__future__ import VideosSearch
 from DragMusic.utils.database import is_on_off
 from DragMusic.utils.formatters import time_to_seconds
+
+SUPPORTED_FORMATS = {"mp4", "mkv", "webm", "avi", "mov"}  # formats that ffmpeg can handle
+
 async def shell_cmd(cmd):
     proc = await asyncio.create_subprocess_shell(
         cmd,
@@ -24,20 +27,21 @@ async def shell_cmd(cmd):
         else:
             return errorz.decode("utf-8")
     return out.decode("utf-8")
+
+
 class YouTubeAPI:
     def __init__(self):
         self.base = "https://www.youtube.com/watch?v="
         self.regex = r"(?:youtube\.com|youtu\.be)"
         self.status = "https://www.youtube.com/oembed?url="
         self.listbase = "https://youtube.com/playlist?list="
-        self.reg = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+        self.reg = re.compile(r"\\x1B(?:[@-Z\\-_]|\\[[0-?]*[ -/]*[@-~])")
+
     async def exists(self, link: str, videoid: Union[bool, str] = None):
         if videoid:
             link = self.base + link
-        if re.search(self.regex, link):
-            return True
-        else:
-            return False
+        return bool(re.search(self.regex, link))
+
     async def url(self, message_1: Message) -> Union[str, None]:
         messages = [message_1]
         if message_1.reply_to_message:
@@ -61,6 +65,7 @@ class YouTubeAPI:
         if offset in (None,):
             return None
         return text[offset : offset + length]
+
     async def details(self, link: str, videoid: Union[bool, str] = None):
         if videoid:
             link = self.base + link
@@ -72,11 +77,9 @@ class YouTubeAPI:
             duration_min = result["duration"]
             thumbnail = result["thumbnails"][0]["url"].split("?")[0]
             vidid = result["id"]
-            if str(duration_min) == "None":
-                duration_sec = 0
-            else:
-                duration_sec = int(time_to_seconds(duration_min))
+            duration_sec = 0 if str(duration_min) == "None" else int(time_to_seconds(duration_min))
         return title, duration_min, duration_sec, thumbnail, vidid
+
     async def title(self, link: str, videoid: Union[bool, str] = None):
         if videoid:
             link = self.base + link
@@ -84,8 +87,8 @@ class YouTubeAPI:
             link = link.split("&")[0]
         results = VideosSearch(link, limit=1)
         for result in (await results.next())["result"]:
-            title = result["title"]
-        return title
+            return result["title"]
+
     async def duration(self, link: str, videoid: Union[bool, str] = None):
         if videoid:
             link = self.base + link
@@ -93,8 +96,8 @@ class YouTubeAPI:
             link = link.split("&")[0]
         results = VideosSearch(link, limit=1)
         for result in (await results.next())["result"]:
-            duration = result["duration"]
-        return duration
+            return result["duration"]
+
     async def thumbnail(self, link: str, videoid: Union[bool, str] = None):
         if videoid:
             link = self.base + link
@@ -102,8 +105,8 @@ class YouTubeAPI:
             link = link.split("&")[0]
         results = VideosSearch(link, limit=1)
         for result in (await results.next())["result"]:
-            thumbnail = result["thumbnails"][0]["url"].split("?")[0]
-        return thumbnail
+            return result["thumbnails"][0]["url"].split("?")[0]
+
     async def video(self, link: str, videoid: Union[bool, str] = None):
         if videoid:
             link = self.base + link
@@ -123,6 +126,7 @@ class YouTubeAPI:
             return 1, stdout.decode().split("\n")[0]
         else:
             return 0, stderr.decode()
+
     async def playlist(self, link, limit, user_id, videoid: Union[bool, str] = None):
         if videoid:
             link = self.listbase + link
@@ -133,12 +137,11 @@ class YouTubeAPI:
         )
         try:
             result = playlist.split("\n")
-            for key in result:
-                if key == "":
-                    result.remove(key)
+            result = [x for x in result if x.strip() != ""]
         except:
             result = []
         return result
+
     async def track(self, link: str, videoid: Union[bool, str] = None):
         if videoid:
             link = self.base + link
@@ -146,19 +149,15 @@ class YouTubeAPI:
             link = link.split("&")[0]
         results = VideosSearch(link, limit=1)
         for result in (await results.next())["result"]:
-            title = result["title"]
-            duration_min = result["duration"]
-            vidid = result["id"]
-            yturl = result["link"]
-            thumbnail = result["thumbnails"][0]["url"].split("?")[0]
-        track_details = {
-            "title": title,
-            "link": yturl,
-            "vidid": vidid,
-            "duration_min": duration_min,
-            "thumb": thumbnail,
-        }
-        return track_details, vidid
+            track_details = {
+                "title": result["title"],
+                "link": result["link"],
+                "vidid": result["id"],
+                "duration_min": result["duration"],
+                "thumb": result["thumbnails"][0]["url"].split("?")[0],
+            }
+        return track_details, result["id"]
+
     async def formats(self, link: str, videoid: Union[bool, str] = None):
         if videoid:
             link = self.base + link
@@ -174,32 +173,23 @@ class YouTubeAPI:
                     str(format["format"])
                 except:
                     continue
-                if not "dash" in str(format["format"]).lower():
+                if "dash" not in str(format["format"]).lower():
                     try:
-                        format["format"]
-                        format["filesize"]
-                        format["format_id"]
-                        format["ext"]
-                        format["format_note"]
+                        formats_available.append(
+                            {
+                                "format": format["format"],
+                                "filesize": format["filesize"],
+                                "format_id": format["format_id"],
+                                "ext": format["ext"],
+                                "format_note": format["format_note"],
+                                "yturl": link,
+                            }
+                        )
                     except:
                         continue
-                    formats_available.append(
-                        {
-                            "format": format["format"],
-                            "filesize": format["filesize"],
-                            "format_id": format["format_id"],
-                            "ext": format["ext"],
-                            "format_note": format["format_note"],
-                            "yturl": link,
-                        }
-                    )
         return formats_available, link
-    async def slider(
-        self,
-        link: str,
-        query_type: int,
-        videoid: Union[bool, str] = None,
-    ):
+
+    async def slider(self, link: str, query_type: int, videoid: Union[bool, str] = None):
         if videoid:
             link = self.base + link
         if "&" in link:
@@ -208,9 +198,9 @@ class YouTubeAPI:
         result = (await a.next()).get("result")
         title = result[query_type]["title"]
         duration_min = result[query_type]["duration"]
-        vidid = result[query_type]["id"]
         thumbnail = result[query_type]["thumbnails"][0]["url"].split("?")[0]
-        return title, duration_min, thumbnail, vidid
+        return title, duration_min, thumbnail, result[query_type]["id"]
+
     async def get_video_info_from_bitflow(self, url: str, video: bool):
         api_url = "https://bitflow.in/api/youtube"
         params = {
@@ -219,107 +209,119 @@ class YouTubeAPI:
             "download": True,
             "api_key": "1spiderkey2"
         }
-        
         async with httpx.AsyncClient() as client:
             response = await client.get(api_url, params=params, timeout=150)
             if response.status_code == 200:
                 return response.json()
             else:
                 return {"status": "error", "message": "Failed to fetch data from Bitflow API."}
-    async def download(
-        self,
-        link: str,
-        mystic,
-        video: Union[bool, str] = None,
-        videoid: Union[bool, str] = None,
-        songaudio: Union[bool, str] = None,
-        songvideo: Union[bool, str] = None,
-        format_id: Union[bool, str] = None,
-        title: Union[bool, str] = None,
-    ) -> str:
+
+    async def download(self, link: str, mystic, video: Union[bool, str] = None,
+                       videoid: Union[bool, str] = None, songaudio: Union[bool, str] = None,
+                       songvideo: Union[bool, str] = None, format_id: Union[bool, str] = None,
+                       title: Union[bool, str] = None) -> str:
         if videoid:
             link = self.base + link
         if "&" in link:
             link = link.split("&")[0]
         loop = asyncio.get_running_loop()
-        bitflow_info = await self.get_video_info_from_bitflow(link, video)
-        def audio_dl(bitflow_info):
-            xyz = os.path.join("downloads", f"{bitflow_info['videoid']}.{bitflow_info['ext']}")
-            ydl_optssx = {
-                "format": "bestaudio/best",
-                "outtmpl": xyz,
-                "geo_bypass": True,
-                "nocheckcertificate": True,
-                "quiet": True,
-                "no_warnings": True,
-            }
-            x = yt_dlp.YoutubeDL(ydl_optssx) 
-            if os.path.exists(xyz):
+        try:
+            bitflow_info = await self.get_video_info_from_bitflow(link, video)
+
+            def audio_dl(bitflow_info):
+                xyz = os.path.join("downloads", f"{bitflow_info['videoid']}.{bitflow_info['ext']}")
+                ydl_optssx = {
+                    "format": "bestaudio/best",
+                    "outtmpl": xyz,
+                    "geo_bypass": True,
+                    "nocheckcertificate": True,
+                    "quiet": True,
+                    "no_warnings": True,
+                }
+                x = yt_dlp.YoutubeDL(ydl_optssx)
+                if os.path.exists(xyz):
+                    return xyz
+                x.download([bitflow_info['url']])
                 return xyz
-            x.download([bitflow_info['url']])
-            return xyz
-        def video_dl(bitflow_info):
-            xyz = os.path.join("downloads", f"{bitflow_info['videoid']}.{bitflow_info['ext']}")
-            ydl_optssx = {
-                "format": "(bestvideo[height<=?720][width<=?1280][ext=mp4])+(bestaudio[ext=m4a])",
-                "outtmpl": xyz,
-                "geo_bypass": True,
-                "nocheckcertificate": True,
-                "quiet": True,
-                "no_warnings": True,
-            }
-            x = yt_dlp.YoutubeDL(ydl_optssx)
-            if os.path.exists(xyz):
+
+            def video_dl(bitflow_info):
+                xyz = os.path.join("downloads", f"{bitflow_info['videoid']}.{bitflow_info['ext']}")
+                ydl_optssx = {
+                    "format": "(bestvideo[height<=?720][width<=?1280][ext=mp4])+(bestaudio[ext=m4a])",
+                    "outtmpl": xyz,
+                    "geo_bypass": True,
+                    "nocheckcertificate": True,
+                    "quiet": True,
+                    "no_warnings": True,
+                }
+                x = yt_dlp.YoutubeDL(ydl_optssx)
+                if os.path.exists(xyz):
+                    return xyz
+                x.download([bitflow_info['url']])
                 return xyz
-            x.download([bitflow_info['url']])
-            return xyz
-        def song_video_dl():
-            formats = f"{format_id}+140"
-            fpath = f"downloads/{title}"
-            ydl_optssx = {
-                "format": formats,
-                "outtmpl": fpath,
-                "geo_bypass": True,
-                "nocheckcertificate": True,
-                "quiet": True,
-                "no_warnings": True,
-                "prefer_ffmpeg": True,
-                "merge_output_format": "mp4",
-            }
-            x = yt_dlp.YoutubeDL(ydl_optssx)
-            x.download([link])
-        def song_audio_dl():
-            fpath = f"downloads/{title}.%(ext)s"
-            ydl_optssx = {
-                "format": format_id,
-                "outtmpl": fpath,
-                "geo_bypass": True,
-                "nocheckcertificate": True,
-                "quiet": True,
-                "no_warnings": True,
-                "prefer_ffmpeg": True,
-                "postprocessors": [
-                    {
-                        "key": "FFmpegExtractAudio",
-                        "preferredcodec": "mp3",
-                        "preferredquality": "192",
-                    }
-                ],
-            }
-            x = yt_dlp.YoutubeDL(ydl_optssx)
-            x.download([link])
-        if songvideo:
-            await loop.run_in_executor(None, song_video_dl)
-            fpath = f"downloads/{title}.mp4"
-            return fpath
-        elif songaudio:
-            await loop.run_in_executor(None, song_audio_dl)
-            fpath = f"downloads/{title}.mp3"
-            return fpath
-        elif video:
-            direct = True
-            downloaded_file = await loop.run_in_executor(None, video_dl, bitflow_info)
-        else:
-            direct = True
-            downloaded_file = await loop.run_in_executor(None, audio_dl, bitflow_info)
-        return downloaded_file, direct  
+
+            def song_video_dl():
+                formats = f"{format_id}+140"
+                fpath = f"downloads/{title}"
+                ydl_optssx = {
+                    "format": formats,
+                    "outtmpl": fpath,
+                    "geo_bypass": True,
+                    "nocheckcertificate": True,
+                    "quiet": True,
+                    "no_warnings": True,
+                    "prefer_ffmpeg": True,
+                    "merge_output_format": "mp4",
+                }
+                x = yt_dlp.YoutubeDL(ydl_optssx)
+                x.download([link])
+
+            def song_audio_dl():
+                fpath = f"downloads/{title}.%(ext)s"
+                ydl_optssx = {
+                    "format": format_id,
+                    "outtmpl": fpath,
+                    "geo_bypass": True,
+                    "nocheckcertificate": True,
+                    "quiet": True,
+                    "no_warnings": True,
+                    "prefer_ffmpeg": True,
+                    "postprocessors": [
+                        {
+                            "key": "FFmpegExtractAudio",
+                            "preferredcodec": "mp3",
+                            "preferredquality": "192",
+                        }
+                    ],
+                }
+                x = yt_dlp.YoutubeDL(ydl_optssx)
+                x.download([link])
+
+            if songvideo:
+                await loop.run_in_executor(None, song_video_dl)
+                fpath = f"downloads/{title}.mp4"
+                return fpath
+            elif songaudio:
+                await loop.run_in_executor(None, song_audio_dl)
+                fpath = f"downloads/{title}.mp3"
+                return fpath
+            elif video:
+                direct = True
+                downloaded_file = await loop.run_in_executor(None, video_dl, bitflow_info)
+            else:
+                direct = True
+                downloaded_file = await loop.run_in_executor(None, audio_dl, bitflow_info)
+
+            # Format check before thumbnail
+            ext = os.path.splitext(downloaded_file)[1].lower().replace(".", "")
+            if ext not in SUPPORTED_FORMATS:
+                print(f"[ERROR] Unsupported format for thumbnail: {ext}")
+                print("[INFO] Skipping thumbnail generation.")
+                return downloaded_file, direct
+
+            return downloaded_file, direct
+
+        except Exception as e:
+            print(f"[ERROR] Failed to download video: {e}")
+            print(traceback.format_exc())
+            return None, False