import asyncio import json import websockets import config import response import ydl_pool type SocketT = websockets.WebSocketServerProtocol async def handler(socket: SocketT) -> None: async def ydl_log_handler(level: response.YdlLogLevel, msg: str) -> None: try: await socket.send(response.ydl_log(level, msg)) except: pass ydls = ydl_pool.Downloader( ydl_pool.YdlLogger( ydl_log_handler, asyncio.get_event_loop(), ) ) async for message in socket: try: data = json.loads(message) print(data) match data['action']: case 'list': # list tracks in album ydls.update_params(data['site'], data.get('proxy', False)) await socket.send(response.playlist( await ydls.get_playlist_items(data['url']), )) case 'download': # download by URL ydls.update_params(data['site'], data.get('proxy', False), data.get('items')) ret = await ydls.download(data['url']) await socket.send(response.ydl_end(ret)) 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.reset_params() ydls.cleanup() async def main() -> None: cfg = config.get() async with websockets.serve(handler, cfg.host, cfg.port): await asyncio.Future() if __name__ == '__main__': asyncio.run(main())