Refactor: merge ydl_wrap with ydl_pool.Downloader, add callbacks

backend/main.py

@@ -6,16 +6,14 @@ import websockets
 import config
 import response
 import ydl_pool
-import ydl_wrap
 
 
 type SocketT = websockets.WebSocketServerProtocol
 
 
-
 async def handler(socket: SocketT) -> None:
 
-    ydls = ydl_pool.Downloaders()
+    ydls = ydl_pool.Downloader(None, None)  # type: ignore  # TODO
 
     async for message in socket:
 
@@ -25,16 +23,14 @@ async def handler(socket: SocketT) -> None:
             match data['action']:
 
                 case 'list':  # list tracks in album
+                    ydls.choose_ydl(data['site'])
                     await socket.send(response.playlist(
-                        await ydl_wrap.get_playlist_items(
+                        await ydls.get_playlist_items(data['url']),
-                            ydls.get_ydl(data['site']),
-                            data['url'],
-                        )
                     ))
 
                 case 'download':  # download by URL
-                    ret = await ydl_wrap.download(
+                    ydls.choose_ydl(data['site'])
-                        ydls.get_ydl(data['site']),
+                    ret = await ydls.download(
                         data['url'],
                         data.get('items'),
                     )

backend/ydl_pool.py

@@ -1,3 +1,6 @@
+import asyncio
+from typing import Callable, Awaitable, Iterable
+
 from yt_dlp import YoutubeDL
 from yt_dlp.postprocessor import FFmpegExtractAudioPP
 
@@ -34,9 +37,12 @@ create_ydl_fn = {
 ydl_fn_keys = create_ydl_fn.keys()
 
 
-class Downloaders:
+class Downloader:
 
-    def __init__(self) -> None:
+    def __init__(
+            self,
+            progress_cb: Callable[[str], Awaitable],
+            lyrics_cb: Callable[[list[str]], Awaitable]) -> None:
 
         self.ydls: dict[str, YoutubeDL | None] = {
             'youtube': None,
@@ -44,10 +50,15 @@ class Downloaders:
             'yandex': None,
         }
 
-    def get_ydl(self, site: str) -> YoutubeDL:
+        self.cur_ydl: YoutubeDL | None = None
+
+        self.progress_cb = progress_cb
+        self.lyrics_cb = lyrics_cb
+
+    def choose_ydl(self, site: str) -> None:
 
-        cfg = config.get()
         ydl = self.ydls[site]
+        cfg = config.get()
 
         if ydl is None:
             ydl = create_ydl_fn[site]()
@@ -60,8 +71,62 @@ class Downloaders:
             if cookies.exists():
                 ydl.params['cookiefile'] = str(cookies)
 
+        self.cur_ydl = ydl
+
+    def get_cur_ydl(self) -> YoutubeDL:
+
+        ydl = self.cur_ydl
+        if ydl is None:
+            raise RuntimeError('ydl object not initialized')
         return ydl
 
+    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(),
+            url,
+        )
+
+    @staticmethod
+    def _target_get_playlist_items(ydl: YoutubeDL, url: str) -> list[str]:
+
+        info = ydl.extract_info(url, download=False, process=True)
+        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:
+
+        return await asyncio.get_event_loop().run_in_executor(
+            None,
+            Downloader._target_download,
+            self.get_cur_ydl(),
+            url,
+            playlist_items,
+        )
+
+    @staticmethod
+    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
+
     def cleanup(self) -> None:
 
         for ydl in self.ydls.values():

backend/ydl_wrap.py

@@ -1,53 +0,0 @@
-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=True)
-    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(
-        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