Fix downloader

bot/handlers/on_chosen/spotify.py

@@ -5,7 +5,7 @@ from aiogram.types import (
 )
 
 from bot.modules.spotify import spotify
-from bot.modules.youtube import youtube
+from bot.modules.youtube import youtube, AgeRestrictedError
 from bot.utils.config import config
 from bot.modules.database import db
 
@@ -16,7 +16,16 @@ router = Router()
 async def on_new_chosen(chosen_result: ChosenInlineResult, bot: Bot):
     song = spotify.songs.from_id(chosen_result.result_id.removeprefix('spot::'))
 
-    bytestream = youtube.songs.search_one(song.full_name).to_bytestream()
+    try:
+        bytestream = await youtube.songs.search_one(song.full_name).to_bytestream()
+    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',
+            reply_markup=None
+        )
+        return
 
     audio = await bot.send_audio(
         chat_id=config.telegram.files_chat,
@@ -26,7 +35,8 @@ async def on_new_chosen(chosen_result: ChosenInlineResult, bot: Bot):
         ),
         thumbnail=URLInputFile(song.thumbnail),
         performer=song.all_artists,
-        title=song.name
+        title=song.name,
+        duration=bytestream.duration,
     )
 
     db.spotify[song.id] = audio.audio.file_id

bot/modules/youtube/__init__.py

@@ -1,7 +1,8 @@
 from .youtube import YouTube
+from pytube.exceptions import AgeRestrictedError
 
 
 youtube = YouTube()
 
 
-__all__ = ['youtube']
+__all__ = ['youtube', 'AgeRestrictedError']

bot/modules/youtube/downloader.py

@@ -4,22 +4,27 @@ from pytube import YouTube, Stream
 from pydub import AudioSegment
 from io import BytesIO
 
+import asyncio
+
 
 @define
 class YouTubeBytestream:
     file: bytes
     filename: str
+    duration: int
 
     @classmethod
     def from_bytestream(
             cls,
             bytestream: BytesIO,
-            filename: str
+            filename: str,
+            duration: float
     ):
         bytestream.seek(0)
         return cls(
             file=bytestream.read(),
-            filename=filename
+            filename=filename,
+            duration=int(duration),
         )
 
     @property
@@ -46,14 +51,22 @@ class Downloader:
             filename=f'{audio_stream.default_filename}.mp3',
         )
 
-    def to_bytestream(self):
+    def __to_bytestream(self):
         audio_io = BytesIO()
         self.audio_stream.stream_to_buffer(audio_io)
         audio_io.seek(0)
 
-        return YouTubeBytestream.from_bytestream(
+        segment = AudioSegment.from_file(
-            AudioSegment.from_file(
             file=audio_io
-            ).export(BytesIO(), format='mp3', codec='libmp3lame'),
+        )
-            self.filename,
+
+        return YouTubeBytestream.from_bytestream(
+            segment.export(BytesIO(), format='mp3', codec='libmp3lame'),
+            self.filename,
+            segment.duration_seconds
+        )
+
+    async def to_bytestream(self):
+        return await asyncio.get_event_loop().run_in_executor(
+            None, self.__to_bytestream
         )

bot/modules/youtube/song.py

@@ -1,7 +1,9 @@
 from attrs import define
 import ytmusicapi
 
-from .downloader import Downloader
+from .downloader import Downloader, YouTubeBytestream
+
+from typing import Awaitable
 
 
 @define
@@ -31,7 +33,7 @@ class SongItem:
     def __str__(self):
         return f"{', '.join(self.artists)} - {self.name}"
 
-    def to_bytestream(self):
+    def to_bytestream(self) -> Awaitable[YouTubeBytestream]:
         return Downloader.from_id(self.id).to_bytestream()