Refactor: merge ydl_wrap with ydl_pool.Downloader, add callbacks
This commit is contained in:
parent
e9c55a67d3
commit
44fc039ac9
3 changed files with 74 additions and 66 deletions
|
@ -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(
|
||||
ydls.get_ydl(data['site']),
|
||||
data['url'],
|
||||
)
|
||||
await ydls.get_playlist_items(data['url']),
|
||||
))
|
||||
|
||||
case 'download': # download by URL
|
||||
ret = await ydl_wrap.download(
|
||||
ydls.get_ydl(data['site']),
|
||||
ydls.choose_ydl(data['site'])
|
||||
ret = await ydls.download(
|
||||
data['url'],
|
||||
data.get('items'),
|
||||
)
|
||||
|
|
|
@ -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():
|
||||
|
|
|
@ -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
|
Loading…
Reference in a new issue