Compare commits

..

2 commits

4 changed files with 88 additions and 76 deletions

View file

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

View file

@ -1,25 +1,29 @@
import json import json
import traceback import traceback
OK = '{"ok":true}' def playlist(items: list[str]) -> str:
def ok_playlist(items: list[str]) -> str:
return json.dumps({ return json.dumps({
"ok": True, "type": "items",
"data": items, "data": items,
}) })
def ok_downloaded(ret: int) -> str: def dl_progress(msg: str) -> str:
return json.dumps({ return json.dumps({
"type": "downloaded", "type": "dl_progress",
"data": msg,
})
def dl_end(ret: int) -> str:
return json.dumps({
"type": "dl_end",
"data": ret, "data": ret,
}) })
def error(ex: Exception) -> str: def error(ex: Exception) -> str:
traceback.print_tb(ex.__traceback__) traceback.print_tb(ex.__traceback__)
return json.dumps({ return json.dumps({
"ok": False, "type": "error",
"error": { "data": {
"type": ex.__class__.__qualname__, "type": ex.__class__.__qualname__,
"message": str(ex), "message": str(ex),
} }

View file

@ -1,3 +1,6 @@
import asyncio
from typing import Callable, Awaitable, Iterable
from yt_dlp import YoutubeDL from yt_dlp import YoutubeDL
from yt_dlp.postprocessor import FFmpegExtractAudioPP from yt_dlp.postprocessor import FFmpegExtractAudioPP
@ -34,9 +37,12 @@ create_ydl_fn = {
ydl_fn_keys = create_ydl_fn.keys() 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] = { self.ydls: dict[str, YoutubeDL | None] = {
'youtube': None, 'youtube': None,
@ -44,10 +50,15 @@ class Downloaders:
'yandex': None, '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] ydl = self.ydls[site]
cfg = config.get()
if ydl is None: if ydl is None:
ydl = create_ydl_fn[site]() ydl = create_ydl_fn[site]()
@ -60,8 +71,62 @@ class Downloaders:
if cookies.exists(): if cookies.exists():
ydl.params['cookiefile'] = str(cookies) 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 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: def cleanup(self) -> None:
for ydl in self.ydls.values(): for ydl in self.ydls.values():

View file

@ -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