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.
This commit is contained in:
parent
7a75805423
commit
9071017dbf
GPG key ID:
0A26CD5B3345D6E3
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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')
|
||||
# <Really-dumb-fix>
|
||||
if artist == 'LINKIN PARK':
|
||||
artist = 'Linkin Park'
|
||||
# </Really-dumb-fix>
|
||||
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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in a new issue