diff --git a/bot/handlers/on_chosen/spotify.py b/bot/handlers/on_chosen/spotify.py index 8f4bdc3..70d08c5 100644 --- a/bot/handlers/on_chosen/spotify.py +++ b/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 diff --git a/bot/modules/youtube/__init__.py b/bot/modules/youtube/__init__.py index 6986775..47e992e 100644 --- a/bot/modules/youtube/__init__.py +++ b/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'] diff --git a/bot/modules/youtube/downloader.py b/bot/modules/youtube/downloader.py index c03a1a7..85a6b98 100644 --- a/bot/modules/youtube/downloader.py +++ b/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( - AudioSegment.from_file( - file=audio_io - ).export(BytesIO(), format='mp3', codec='libmp3lame'), - self.filename, + segment = AudioSegment.from_file( + file=audio_io + ) + + 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 ) diff --git a/bot/modules/youtube/song.py b/bot/modules/youtube/song.py index 6704a7a..c3dedfa 100644 --- a/bot/modules/youtube/song.py +++ b/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()