diff --git a/backend/id3pp.py b/backend/id3pp.py index 960d90d..ae9f657 100644 --- a/backend/id3pp.py +++ b/backend/id3pp.py @@ -20,6 +20,8 @@ class InfoYouTubePP(PostProcessor): information['track'] = information['title'] if not 'artist' in information: information['artist'] = information['channel'].removesuffix(' - Topic') + if not 'artists' in information: + information['artists'] = [information['artist']] if not 'album' in information: try: @@ -47,10 +49,10 @@ class ID3TagsPP(PostProcessor): file = mp3.MP3(information['filepath']) title = information['track'] - artist = information['artist'] + artists: list[str] = information['artists'] file['TIT2'] = id3.TIT2(encoding=ENC_UTF8, text=title) - file['TPE1'] = id3.TPE1(encoding=ENC_UTF8, text=artist) + file['TPE1'] = id3.TPE1(encoding=ENC_UTF8, text=artists) # multiple values are null-separated if 'album' in information: file['TALB'] = id3.TALB(encoding=ENC_UTF8, text=information['album']) if 'release_year' in information: @@ -61,7 +63,7 @@ class ID3TagsPP(PostProcessor): file['TCON'] = id3.TCON(encoding=ENC_UTF8, text=information['genre']) try: - lyr_url = genius.search(title, artist) + lyr_url = genius.search(title, artists[0]) file['USLT'] = id3.USLT(encoding=ENC_UTF8, text=genius.parse(lyr_url)) except: pass diff --git a/backend/main.py b/backend/main.py index 2d722e4..c954104 100644 --- a/backend/main.py +++ b/backend/main.py @@ -1,6 +1,5 @@ import asyncio import json -from typing import Any import secrets import time @@ -8,48 +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: - 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: pass command to the thread started in `init` - await socket.send(response.OK) - case _: - await socket.send(response.error_field('action')) + + ydls = ydl_pool.Downloaders() + + async for message in socket: + + 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/test_pp.py b/backend/test_pp.py index 0c95a13..6ad0205 100644 --- a/backend/test_pp.py +++ b/backend/test_pp.py @@ -25,6 +25,7 @@ INFO_AFTER = { **INFO_BEFORE, "track": "A Line in the Sand", "artist": "Linkin Park", + "artists": ["Linkin Park"], "album": "The Hunting Party", "track_number": 12, } @@ -39,7 +40,7 @@ class TestPostProcessorsOnFakeData(TestCase): def test_infoytpp(self) -> None: _, info = id3pp.InfoYouTubePP().run(INFO_BEFORE) - self.assertEqual(info, INFO_AFTER) + self.assertDictEqual(info, INFO_AFTER) def test_id3(self) -> None: 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