Working WS server, yt-dlp objects "pool", async wrappers

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:

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),
+        }
     })

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()

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