From 2b47b002e68f06e02c70d2dc99c514b673b14f1d Mon Sep 17 00:00:00 2001 From: DarkCat09 Date: Fri, 3 May 2024 18:42:53 +0400 Subject: [PATCH] Working WS server, yt-dlp objects "pool", async wrappers --- backend/main.py | 73 ++++++++++++++++++++++++++------------------- backend/response.py | 22 ++++++++------ backend/ydl_pool.py | 61 +++++++++++++++++++++++++++++++++++++ backend/ydl_wrap.py | 33 ++++++++++++++++++++ 4 files changed, 150 insertions(+), 39 deletions(-) create mode 100644 backend/ydl_pool.py create mode 100644 backend/ydl_wrap.py diff --git a/backend/main.py b/backend/main.py index 01b81a1..c954104 100644 --- a/backend/main.py +++ b/backend/main.py @@ -7,49 +7,62 @@ import time import websockets import response +import ydl_pool +import ydl_wrap # TODO: config HOST = '127.0.0.1' PORT = 5678 -# TODO: make dict with functions creating YoutubeDL for each site instead of this set -SITES = {'youtube', 'yandex'} - -sessions = {} - def generate_key() -> str: + return hex(time.time_ns()) + secrets.token_hex(2) async def handler(socket: websockets.WebSocketServerProtocol) -> None: + + ydls = ydl_pool.Downloaders() + async for message in socket: - data = json.loads(await socket.recv()) - if 'action' not in data: - await socket.send(response.error_field('action')) - return - match data['action']: - case 'init': # create session - if data.get('site') not in SITES: - await socket.send(response.error_field('site')) - return - key = generate_key() - sessions[key] = "" # TODO - await socket.send(response.ok_init(key)) - case 'list': # list tracks in album - if 'url' not in data: - await socket.send(response.error_field('url')) - return - await socket.send(response.ok_playlist(['title 1', 'title 2'])) # TODO - case 'download': # download by URL - if 'url' not in data: - await socket.send(response.error_field('url')) - return - # TODO: start thread and send all yt-dlp's output to client - await socket.send(response.OK) - case _: - await socket.send(response.error_field('action')) + + try: + data = json.loads(message) + + match data['action']: + + case 'list': # list tracks in album + await socket.send(response.ok_playlist( + await ydl_wrap.get_playlist_items( + ydls.get_ydl(data['site']), + data['url'], + ) + )) + + case 'download': # download by URL + # TODO: send all yt-dlp's output to client + await socket.send(response.ok_download( + await ydl_wrap.download( + ydls.get_ydl(data['site']), + data['url'], + data.get('items'), + ) + )) + + # TODO: cancellation + + case _: + raise ValueError('invalid "action" field value') + + except websockets.ConnectionClosed: + break + + except Exception as ex: + if not socket.closed: + await socket.send(response.error(ex)) + + ydls.cleanup() async def main() -> None: diff --git a/backend/response.py b/backend/response.py index 63173af..7ea223c 100644 --- a/backend/response.py +++ b/backend/response.py @@ -1,22 +1,26 @@ import json +import traceback OK = '{"ok":true}' -def ok_init(key: str) -> str: - return json.dumps({ - "ok": True, - "data": key, - }) - def ok_playlist(items: list[str]) -> str: return json.dumps({ "ok": True, "data": items, }) -def error_field(field: str) -> str: +def ok_download(ret: int) -> str: + return json.dumps({ + "ok": True, + "data": ret, + }) + +def error(ex: Exception) -> str: + traceback.print_tb(ex.__traceback__) return json.dumps({ "ok": False, - "error": "field", - "data": field, + "error": { + "type": ex.__class__.__qualname__, + "message": str(ex), + } }) diff --git a/backend/ydl_pool.py b/backend/ydl_pool.py new file mode 100644 index 0000000..b083517 --- /dev/null +++ b/backend/ydl_pool.py @@ -0,0 +1,61 @@ +from yt_dlp import YoutubeDL +from yt_dlp.postprocessor import FFmpegExtractAudioPP + +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() + ydl.params['proxy'] = 'http://127.0.0.1:1080' # TODO + return ydl + + @staticmethod + def yandex() -> YoutubeDL: + return YoutubeDL() # TODO: cookies + + +create_ydl_fn = { + 'youtube': _CreateYDL.youtube, + 'yt_proxied': _CreateYDL.yt_proxied, + 'yandex': _CreateYDL.yandex, +} + +ydl_fn_keys = create_ydl_fn.keys() + + +class Downloaders: + + def __init__(self) -> None: + + self.ydls: dict[str, YoutubeDL | None] = { + 'youtube': None, + 'yt_proxied': None, + 'yandex': None, + } + + def get_ydl(self, site: str) -> YoutubeDL: + + ydl = self.ydls[site] + if ydl is None: + ydl = create_ydl_fn[site]() + ydl.params['trim_file_name'] = 40 # TODO: config + # artists.0 instead of artist, because it can contain "feat. ..." + ydl.params['outtmpl']['default'] = 'music/%(artists.0)s/%(album)s/%(track)s.%(ext)s' + ydl.add_post_processor(id3pp.ID3TagsPP(), when='post_process') + return ydl + + def cleanup(self) -> None: + + for ydl in self.ydls.values(): + if ydl is not None: + ydl.close() diff --git a/backend/ydl_wrap.py b/backend/ydl_wrap.py new file mode 100644 index 0000000..560c8cf --- /dev/null +++ b/backend/ydl_wrap.py @@ -0,0 +1,33 @@ +import asyncio +from typing import Iterable + +from yt_dlp import YoutubeDL + + +async def get_playlist_items(ydl: YoutubeDL, url: str) -> list[str]: + + return await asyncio.get_event_loop().run_in_executor(None, _target_get_playlist_items, ydl, url) + + +def _target_get_playlist_items(ydl: YoutubeDL, url: str) -> list[str]: + + info = ydl.extract_info(url, download=False, process=False) + if info is None: + raise RuntimeError('ydl.extract_info returned None') + return [entry['title'] for entry in info['entries']] + + +async def download(ydl: YoutubeDL, url: str, playlist_items: Iterable[int] | None = None) -> int: + + return await asyncio.get_event_loop().run_in_executor(None, _target_download, ydl, url, playlist_items) + + +def _target_download(ydl: YoutubeDL, url: str, playlist_items: Iterable[int] | None = None) -> int: + + if playlist_items: + ydl.params['playlist_items'] = ','.join(str(i) for i in playlist_items) + + ret = ydl.download(url) + del ydl.params['playlist_items'] + + return ret