From 9071017dbf3024271f82f9f1d085ac44b0782bd9 Mon Sep 17 00:00:00 2001 From: DarkCat09 Date: Tue, 28 May 2024 13:21:23 +0400 Subject: [PATCH] Refactor: unify YDL object and proxy cfg field This would allow to add sites much simplier or even drop the ydl_fn_keys. Main purpose of the refactoring is a code cleanup. --- backend/config.py | 7 +-- backend/id3pp.py | 21 ++++--- backend/main.py | 12 ++-- backend/test_pp.py | 4 +- backend/ydl_pool.py | 137 +++++++++++++++----------------------------- 5 files changed, 68 insertions(+), 113 deletions(-) diff --git a/backend/config.py b/backend/config.py index 35dfbb3..1f95a8e 100644 --- a/backend/config.py +++ b/backend/config.py @@ -20,11 +20,8 @@ class Config: os.getenv('TRACK_FILE_TMPL') or '%(track)S.%(ext)s', ) - # Proxy URL for yt_proxied downloader (can be used for geo-restricted content) - self.yt_proxy = os.getenv('YT_PROXY') or None - - # Proxy for soundcloud - self.sc_proxy = os.getenv('SC_PROXY') or None + # `or None` -- defaults to None also if PROXY is an empty string + self.proxy = os.getenv('PROXY') or None self.save_lyrics = _parse_bool(os.getenv('SAVE_LYRICS'), True) self.save_cover = _parse_bool(os.getenv('SAVE_COVER'), True) diff --git a/backend/id3pp.py b/backend/id3pp.py index 4348568..323d785 100644 --- a/backend/id3pp.py +++ b/backend/id3pp.py @@ -10,19 +10,24 @@ import genius ENC_UTF8 = 3 -class InfoYouTubePP(PostProcessor): - '''Generates YT Music fields in info if necessary. +class InfoGenPP(PostProcessor): + '''Generates track, artist(s), album, track_number fields + from other info fields if they are not present in current info dict. Must be run before downloading and post-processing with - FFmpegExtractAudioPP and ID3YouTubePP, so use only with - when<="before_dl" ("pre_process" also suits, see yt_dlp.utils.POSTPROCESS_WHEN) - for correct file path and ID3 tags''' + FFmpegExtractAudioPP and ID3TagsPP, so use only with + when="pre_process" for correct file path and ID3 tags''' def run(self, information): if not 'track' in information: information['track'] = information['title'] if not 'artist' in information: - information['artist'] = information['uploader'].removesuffix(' - Topic') + artist = information['uploader'].removesuffix(' - Topic') + # + if artist == 'LINKIN PARK': + artist = 'Linkin Park' + # + information['artist'] = artist if not 'artists' in information: information['artists'] = [information['artist']] @@ -44,8 +49,8 @@ class InfoYouTubePP(PostProcessor): class ID3TagsPP(PostProcessor): - '''Inserts ID3 tags after all PPs (for YT: InfoYouTubePP and FFmpegExtractAudioPP), - triggers searching and parsing lyrics from Genius''' + '''Inserts ID3 tags including lyrics (see parser in backend/genius.py). + Must be run after all PPs ()''' def __init__(self) -> None: self.cfg = config.get() diff --git a/backend/main.py b/backend/main.py index 7c56d21..691b8e0 100644 --- a/backend/main.py +++ b/backend/main.py @@ -30,21 +30,19 @@ async def handler(socket: SocketT) -> None: try: data = json.loads(message) + print(data) match data['action']: case 'list': # list tracks in album - ydls.choose_ydl(data['site']) + ydls.update_params(data['site'], data.get('proxy', False)) await socket.send(response.playlist( await ydls.get_playlist_items(data['url']), )) case 'download': # download by URL - ydls.choose_ydl(data['site']) - ret = await ydls.download( - data['url'], - data.get('items'), - ) + ydls.update_params(data['site'], data.get('proxy', False), data.get('items')) + ret = await ydls.download(data['url']) await socket.send(response.ydl_end(ret)) case _: @@ -57,6 +55,8 @@ async def handler(socket: SocketT) -> None: if not socket.closed: await socket.send(response.error(ex)) + ydls.reset_params() + ydls.cleanup() diff --git a/backend/test_pp.py b/backend/test_pp.py index 3cfa0fd..16d2f14 100644 --- a/backend/test_pp.py +++ b/backend/test_pp.py @@ -15,7 +15,7 @@ TEST_MP3 = 'music/test.mp3' INFO_BEFORE = { "filepath": TEST_MP3, "title": "A Line in the Sand", - "channel": "Linkin Park - Topic", + "uploader": "Linkin Park - Topic", "playlist": "Album \u2013 The Hunting Party", "playlist_index": 12, "release_year": 2015, @@ -44,7 +44,7 @@ class TestPostProcessorsOnFakeData(TestCase): def test_infoytpp(self) -> None: - _, info = id3pp.InfoYouTubePP().run(INFO_BEFORE) + _, info = id3pp.InfoGenPP().run(INFO_BEFORE) self.assertDictEqual(info, INFO_AFTER) def test_id3(self) -> None: diff --git a/backend/ydl_pool.py b/backend/ydl_pool.py index 825e532..a82b136 100644 --- a/backend/ydl_pool.py +++ b/backend/ydl_pool.py @@ -10,46 +10,7 @@ import response import id3pp -class _CreateYDL: - - @staticmethod - def youtube() -> YoutubeDL: - ydl = YoutubeDL({'format': 'ba'}) - ydl.add_post_processor(id3pp.InfoYouTubePP(), when='before_dl') - ydl.add_post_processor(FFmpegExtractAudioPP(preferredcodec='mp3'), when='post_process') - return ydl - - @staticmethod - def yt_proxied() -> YoutubeDL: - ydl = _CreateYDL.youtube() - proxy = config.get().yt_proxy - if proxy is not None: - ydl.params['proxy'] = proxy - return ydl - - @staticmethod - def soundcloud() -> YoutubeDL: - ydl = YoutubeDL({'format': 'ba[ext=mp3]/ba'}) - ydl.add_post_processor(id3pp.InfoYouTubePP(), when='pre_process') - proxy = config.get().sc_proxy - if proxy is not None: - ydl.params['proxy'] = proxy - ydl.add_post_processor(FFmpegExtractAudioPP(preferredcodec='mp3'), when='post_process') - return ydl - - @staticmethod - def yandex() -> YoutubeDL: - return YoutubeDL() - - -create_ydl_fn = { - 'youtube': _CreateYDL.youtube, - 'yt_proxied': _CreateYDL.yt_proxied, - 'soundcloud': _CreateYDL.soundcloud, - 'yandex': _CreateYDL.yandex, -} - -ydl_fn_keys = create_ydl_fn.keys() +ydl_fn_keys = {'youtube', 'yandex', 'soundcloud'} # need process=True for track title in extract_info output NP_YDLS = {'yandex', 'soundcloud'} @@ -83,98 +44,90 @@ class Downloader: def __init__(self, logger: YdlLogger | None = None) -> None: - self.ydls: dict[str, YoutubeDL | None] = { - key: None - for key in ydl_fn_keys - } + self.cfg = config.get() - self.cur_ydl: YoutubeDL | None = None - self.cur_site = '' + self.ydl = YoutubeDL({'format': 'ba[ext=mp3]/ba/b'}) + self.ydl.add_post_processor(id3pp.InfoGenPP(), when='pre_process') + # Note: it skips converting if downloaded file is already MP3 + self.ydl.add_post_processor(FFmpegExtractAudioPP(preferredcodec='mp3'), when='post_process') + # ID3 tags are set after all pre/post-processings + self.ydl.add_post_processor(id3pp.ID3TagsPP(), when='post_process') - self.id3pp_obj = id3pp.ID3TagsPP() + if logger is not None: + self.ydl.params['logger'] = logger - self.logger = logger + self.ydl.params['outtmpl']['default'] = self.cfg.tmpl - def choose_ydl(self, site: str) -> None: + self.need_process = False + self.updated_params = [] - ydl = self.ydls[site] - cfg = config.get() + def update_params( + self, + site: str, + proxy: bool = False, + items: Iterable[int] | None = None) -> None: - if ydl is None: - ydl = create_ydl_fn[site]() + self.need_process = site in NP_YDLS - if self.logger is not None: - ydl.params['logger'] = self.logger + cookies = self.cfg.cookies_dir / (site + '.txt') + if cookies.exists(): + self.ydl.params['cookiefile'] = str(cookies) + self.updated_params.append('cookiefile') - ydl.params['outtmpl']['default'] = cfg.tmpl - ydl.add_post_processor(self.id3pp_obj, when='post_process') + if proxy and self.cfg.proxy is not None: + self.ydl.params['proxy'] = self.cfg.proxy + self.updated_params.append('proxy') - cookies = cfg.cookies_dir / (site + '.txt') - if cookies.exists(): - ydl.params['cookiefile'] = str(cookies) + if items: + self.ydl.params['playlist_items'] = ( + ','.join(str(i) for i in items) + ) + self.updated_params.append('playlist_items') - self.cur_ydl = ydl - self.cur_site = site + def reset_params(self) -> None: - def get_cur_ydl(self) -> YoutubeDL: - - ydl = self.cur_ydl - if ydl is None: - raise RuntimeError('ydl object not initialized') - return ydl + for param in self.updated_params: + del self.ydl.params[param] + self.updated_params.clear() + self.need_process = False async def get_playlist_items(self, url: str) -> list[str]: return await asyncio.get_event_loop().run_in_executor( None, Downloader._target_get_playlist_items, - self.get_cur_ydl(), + self.ydl, url, - self.cur_site in NP_YDLS, + self.need_process, ) @staticmethod def _target_get_playlist_items(ydl: YoutubeDL, url: str, process: bool) -> list[str]: info = ydl.extract_info(url, download=False, process=process) + if info is None: raise RuntimeError('ydl.extract_info returned None') + return [ entry['track'] if 'track' in entry else entry['title'] for entry in info['entries'] ] - async def download( - self, - url: str, - playlist_items: Iterable[int] | None = None) -> int: + async def download(self, url: str) -> int: return await asyncio.get_event_loop().run_in_executor( None, Downloader._target_download, - self.get_cur_ydl(), + self.ydl, url, - playlist_items, ) @staticmethod - def _target_download( - ydl: YoutubeDL, - url: str, - playlist_items: Iterable[int] | None = None) -> int: + def _target_download(ydl: YoutubeDL, url: str) -> int: - if playlist_items: - ydl.params['playlist_items'] = ','.join(str(i) for i in playlist_items) - - ret = ydl.download(url) - - if playlist_items: - del ydl.params['playlist_items'] - - return ret + return ydl.download(url) def cleanup(self) -> None: - for ydl in self.ydls.values(): - if ydl is not None: - ydl.close() + self.ydl.close()