musicdlp/backend/ydl_pool.py
DarkCat09 9071017dbf
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.
2024-05-28 13:21:23 +04:00

133 lines
3.8 KiB
Python

import asyncio
import logging
from typing import Callable, Awaitable, Iterable
from yt_dlp import YoutubeDL
from yt_dlp.postprocessor import FFmpegExtractAudioPP
import config
import response
import id3pp
ydl_fn_keys = {'youtube', 'yandex', 'soundcloud'}
# need process=True for track title in extract_info output
NP_YDLS = {'yandex', 'soundcloud'}
class YdlLogger:
def __init__(
self,
log_cb: Callable[[response.YdlLogLevel, str], Awaitable],
loop: asyncio.AbstractEventLoop) -> None:
self.log_cb = log_cb
self.loop = loop
self.mdlp_logger = logging.getLogger('musicdlp')
def debug(self, msg: str) -> None:
asyncio.run_coroutine_threadsafe(self.log_cb('debug', msg), self.loop)
def info(self, _: str) -> None: # afaik not used in yt-dlp
pass
def warning(self, msg: str) -> None:
asyncio.run_coroutine_threadsafe(self.log_cb('warning', msg), self.loop)
def error(self, msg: str) -> None:
self.mdlp_logger.error(msg)
asyncio.run_coroutine_threadsafe(self.log_cb('error', msg), self.loop)
class Downloader:
def __init__(self, logger: YdlLogger | None = None) -> None:
self.cfg = config.get()
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')
if logger is not None:
self.ydl.params['logger'] = logger
self.ydl.params['outtmpl']['default'] = self.cfg.tmpl
self.need_process = False
self.updated_params = []
def update_params(
self,
site: str,
proxy: bool = False,
items: Iterable[int] | None = None) -> None:
self.need_process = site in NP_YDLS
cookies = self.cfg.cookies_dir / (site + '.txt')
if cookies.exists():
self.ydl.params['cookiefile'] = str(cookies)
self.updated_params.append('cookiefile')
if proxy and self.cfg.proxy is not None:
self.ydl.params['proxy'] = self.cfg.proxy
self.updated_params.append('proxy')
if items:
self.ydl.params['playlist_items'] = (
','.join(str(i) for i in items)
)
self.updated_params.append('playlist_items')
def reset_params(self) -> None:
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.ydl,
url,
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) -> int:
return await asyncio.get_event_loop().run_in_executor(
None,
Downloader._target_download,
self.ydl,
url,
)
@staticmethod
def _target_download(ydl: YoutubeDL, url: str) -> int:
return ydl.download(url)
def cleanup(self) -> None:
self.ydl.close()