Fix youtube downloading wrong track, saving exceptions, attempt to fix deezer

bot/__init__.py

@@ -6,7 +6,9 @@ async def runner():
     from .common import dp, bot
 
     from . import handlers, callbacks
+    from .modules.error import on_error
 
+    dp.error.register(on_error)
     dp.include_routers(
         handlers.router,
         callbacks.router,
@@ -16,14 +18,20 @@ async def runner():
     await dp.start_polling(bot)
 
 
+def plugins():
+    import nest_asyncio
+    from rich import traceback
+    from icecream import ic
+
+    nest_asyncio.apply()
+    traceback.install()
+    ic.configureOutput(includeContext=True)
+
+
 def main():
     import asyncio
 
-    from rich.traceback import install
-    install(show_locals=True)
-
-    from nest_asyncio import apply
-    apply()
+    plugins()
 
     print('Starting...')
     with contextlib.suppress(KeyboardInterrupt):

bot/common.py

@@ -1,8 +1,11 @@
 from aiogram import Bot, Dispatcher
 from bot.modules.fsm import InDbStorage
+from rich.console import Console
 from .utils.config import config
 
 bot = Bot(token=config.telegram.bot_token)
 dp = Dispatcher(storage=InDbStorage())
+console = Console()
 
-__all__ = ['bot', 'dp', 'config']
+
+__all__ = ['bot', 'dp', 'config', 'console']

bot/handlers/on_chosen/spotify.py

@@ -6,45 +6,108 @@ from aiogram.types import (
 
 from bot.modules.spotify import spotify
 from bot.modules.youtube import youtube, AgeRestrictedError
+from bot.modules.youtube.song import SongItem
+from bot.modules.deezer import deezer
 from bot.utils.config import config
 from bot.modules.database import db
 
 router = Router()
 
 
+def not_strict_name(song, yt_song):
+    if 'feat' in yt_song.name.lower():
+        return any(artist.lower() in yt_song.name.lower() for artist in song.artists)
+    else:
+        return False
+
+
 @router.chosen_inline_result(F.result_id.startswith('spot::'))
 async def on_new_chosen(chosen_result: ChosenInlineResult, bot: Bot):
     song = spotify.songs.from_id(chosen_result.result_id.removeprefix('spot::'))
 
+    bytestream = None
+    audio = None
+
+    yt_song: SongItem = youtube.songs.search_one(
+        song.full_name,
+        exact_match=True,
+    )
+    if ((song.all_artists != yt_song.all_artists or song.name != yt_song.name)
+            and not not_strict_name(song, yt_song)):
+        await bot.edit_message_caption(
+            inline_message_id=chosen_result.inline_message_id,
+            caption='🙄 Cannot find this song on YouTube, trying Deezer...',
+            reply_markup=None,
+            parse_mode='HTML',
+        )
+        bytestream = False
+
     try:
-        bytestream = await youtube.songs.search_one(song.full_name).to_bytestream()
+        if bytestream is None:
+            bytestream = await yt_song.to_bytestream()
+
+            audio = await bot.send_audio(
+                chat_id=config.telegram.files_chat,
+                audio=BufferedInputFile(
+                    file=bytestream.file,
+                    filename=bytestream.filename,
+                ),
+                thumbnail=URLInputFile(song.thumbnail),
+                performer=song.all_artists,
+                title=song.name,
+                duration=bytestream.duration,
+            )
+            db.youtube[yt_song.id] = audio.audio.file_id
+
     except AgeRestrictedError:
         await bot.edit_message_caption(
             inline_message_id=chosen_result.inline_message_id,
-            caption='🔞 This song is age restricted, so I can\'t download it. '
-                    'Try downloading it from Deezer or SoundCloud',
+            caption='🔞 This song is age restricted, trying Deezer...',
+            reply_markup=None,
+            parse_mode='HTML',
+        )
+
+    if not bytestream:
+        try:
+            deezer_song = await deezer.songs.search_one(
+                song.full_name,
+            )
+
+            bytestream = await (
+                await deezer.downloader.from_id(deezer_song.id)
+            ).to_bytestream()
+
+            audio = await bot.send_audio(
+                chat_id=config.telegram.files_chat,
+                audio=BufferedInputFile(
+                    file=bytestream.file,
+                    filename=bytestream.filename,
+                ),
+                thumbnail=URLInputFile(bytestream.song.thumbnail),
+                performer=bytestream.song.all_artists,
+                title=bytestream.song.name,
+                duration=bytestream.song.duration,
+            )
+
+            db.deezer[bytestream.song.id] = audio.audio.file_id
+        except Exception as e:
+            assert e
+
+    if audio:
+        db.spotify[song.id] = audio.audio.file_id
+
+        await bot.edit_message_media(
+            inline_message_id=chosen_result.inline_message_id,
+            media=InputMediaAudio(media=audio.audio.file_id),
             reply_markup=None
         )
-        return
 
-    audio = await bot.send_audio(
-        chat_id=config.telegram.files_chat,
-        audio=BufferedInputFile(
-            file=bytestream.file,
-            filename=bytestream.filename,
-        ),
-        thumbnail=URLInputFile(song.thumbnail),
-        performer=song.all_artists,
-        title=song.name,
-        duration=bytestream.duration,
-    )
-
-    db.spotify[song.id] = audio.audio.file_id
-
-    await bot.edit_message_media(
-        inline_message_id=chosen_result.inline_message_id,
-        media=InputMediaAudio(media=audio.audio.file_id),
-        reply_markup=None
-    )
+    else:
+        await bot.edit_message_caption(
+            inline_message_id=chosen_result.inline_message_id,
+            caption='🤷‍♂️ Cannot download this song',
+            reply_markup=None,
+            parse_mode='HTML',
+        )
 
     await db.occasionally_write()

bot/keyboards/inline/search_variants.py

@@ -0,0 +1,32 @@
+from aiogram.utils.keyboard import (InlineKeyboardMarkup, InlineKeyboardButton,
+                                    InlineKeyboardBuilder)
+
+
+deezer = {
+    'd': '🎵 Search in Deezer'
+}
+soundcloud = {
+    'c': '☁️ Search in SoundCloud'
+}
+youtube = {
+    'y': '▶️ Search in YouTube'
+}
+spotify = {
+    's': '🎧 Search in Spotify'
+}
+
+
+def get_search_variants_kb(
+        query: str,
+        services: dict[str, str],
+) -> InlineKeyboardMarkup:
+    buttons = [
+        [
+            InlineKeyboardButton(
+                text=services[key],
+                switch_inline_query_current_chat=f'{key}:{query}'
+            )
+        ] for key in services.keys()
+    ]
+
+    return InlineKeyboardBuilder(buttons).as_markup()

bot/modules/database/db.py

@@ -15,6 +15,7 @@ class Db(object):
         self.fsm = DBDict('fsm')
         self.config = DBDict('config')
         self.inline = DBDict('inline')
+        self.errors = DBDict('errors')
         self.spotify = DBDict('spotify')
         self.deezer = DBDict('deezer')
         self.youtube = DBDict('youtube')

bot/modules/deezer/downloader.py

@@ -44,12 +44,18 @@ class Downloader:
             driver: DeezerDriver
     ):
         track = await driver.reverse_get_track(song_id)
-        return cls(
-            song_id=str(song_id),
-            driver=driver,
-            track=track['results'],
-            song=await FullSongItem.from_deezer(track)
-        )
+        try:
+            return cls(
+                song_id=str(song_id),
+                driver=driver,
+                track=track['results'],
+                song=await FullSongItem.from_deezer(track)
+            )
+        except KeyError:
+            from icecream import ic
+            ic(track)
+            await driver.renew_engine()
+            return await cls.build(song_id, driver)
 
     async def to_bytestream(self) -> DeezerBytestream:
         quality = track_formats.MP3_128

bot/modules/deezer/driver.py

@@ -34,3 +34,6 @@ class DeezerDriver:
         )
 
         return data['data']
+
+    async def renew_engine(self):
+        self.engine = await self.engine.from_arl(self.engine.arl)

bot/modules/deezer/engine.py

@@ -20,6 +20,7 @@ HTTP_HEADERS = {
 @define
 class DeezerEngine:
     cookies: dict
+    arl: str = None
     token: str = None
 
     @classmethod
@@ -34,6 +35,7 @@ class DeezerEngine:
 
         return cls(
             cookies=cookies,
+            arl=arl,
             token=token
         )
 

bot/modules/error/__init__.py

@@ -0,0 +1 @@
+from .handler import on_error

bot/modules/error/handler.py

@@ -0,0 +1,53 @@
+from bot.common import console
+from aiogram.types.error_event import ErrorEvent
+from aiogram import Bot
+
+from rich.traceback import Traceback
+
+from bot.modules.database import db
+
+from dataclasses import dataclass
+
+
+@dataclass
+class Error:
+    traceback: Traceback
+    inline_message_id: str | None = None
+
+
+async def on_error(event: ErrorEvent, bot: Bot):
+    import os
+    import base64
+
+    error_id = base64.urlsafe_b64encode(os.urandom(6)).decode()
+
+    traceback = Traceback.from_exception(
+        type(event.exception),
+        event.exception,
+        event.exception.__traceback__,
+        show_locals=True,
+        max_frames=1,
+    )
+
+    if event.update.chosen_inline_result:
+        db.errors[error_id] = Error(
+            traceback=traceback,
+            inline_message_id=event.update.chosen_inline_result.inline_message_id,
+        )
+
+        await bot.edit_message_caption(
+            inline_message_id=event.update.chosen_inline_result.inline_message_id,
+            caption=f'💔 <b>ERROR</b> occurred. Use this code to get more information: '
+                    f'<code>{error_id}</code>',
+            parse_mode='HTML',
+        )
+
+    else:
+        db.errors[error_id] = Error(
+            traceback=traceback,
+        )
+
+    console.print(f'[red]{error_id} occurred[/]')
+    console.print(event)
+    console.print(traceback)
+    console.print(f'-{error_id}-')

bot/modules/error/pretty.py

@@ -0,0 +1,50 @@
+import os
+import traceback
+import contextlib
+import re
+
+
+class PrettyException:
+    def __init__(self, e: Exception):
+        self.pretty_exception = f"""
+❌ Error! Report it to admins:
+🐊 <code>{e.__traceback__.tb_frame.f_code.co_filename.replace(os.getcwd(), "")}\r
+</code>:{e.__traceback__.tb_frame.f_lineno}
+😍 {e.__class__.__name__}
+👉 {"".join(traceback.format_exception_only(e)).strip()}
+
+⬇️ Trace:
+{self.get_full_stack()}
+"""
+
+    @staticmethod
+    def get_full_stack():
+        full_stack = traceback.format_exc().replace(
+            "Traceback (most recent call last):\n", ""
+        )
+
+        line_regex = r'  File "(.*?)", line ([0-9]+), in (.+)'
+
+        def format_line(line: str) -> str:
+            filename_, lineno_, name_ = re.search(line_regex, line).groups()
+            with contextlib.suppress(Exception):
+                filename_ = os.path.basename(filename_)
+
+            return (
+                f"🤯 <code>{filename_}:{lineno_}</code> (<b>in</b>"
+                f" <code>{name_}</code> call)"
+            )
+
+        full_stack = "\n".join(
+            [
+                format_line(line)
+                if re.search(line_regex, line)
+                else f"<code>{line}</code>"
+                for line in full_stack.splitlines()
+            ]
+        )
+
+        return full_stack
+
+    def __str__(self):
+        return self.pretty_exception

bot/modules/youtube/song.py

@@ -38,16 +38,26 @@ class SongItem(BaseSongItem):
 class Songs(object):
     ytm: ytmusicapi.YTMusic
 
-    def search(self, query: str, limit: int = 10) -> list[SongItem] | None:
-        r = self.ytm.search(query, limit=limit, filter='songs')
+    def search(
+            self,
+            query: str,
+            limit: int = 10,
+            exact_match: bool = False
+    ) -> list[SongItem] | None:
+        r = self.ytm.search(
+            query,
+            limit=limit,
+            filter='songs',
+            ignore_spelling=exact_match
+        )
 
         if r is None:
             return None
 
         return [SongItem.from_youtube(song_item) for song_item in r]
 
-    def search_one(self, query: str) -> SongItem | None:
-        return (self.search(query, limit=1) or [None])[0]
+    def search_one(self, query: str, exact_match: bool = False) -> SongItem | None:
+        return (self.search(query, limit=1, exact_match=exact_match) or [None])[0]
 
     def from_id(self, song_id: str) -> SongItem | None:
         r = self.ytm.get_song(song_id)

pyproject.toml

@@ -20,6 +20,7 @@ pytube = "^15.0.0"
 pydub = "^0.25.1"
 aiohttp = "^3.8.6"
 nest-asyncio = "^1.5.8"
+icecream = "^2.1.3"
 
 
 [build-system]