import asyncio import json 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() async for message in socket: try: data = json.loads(message) match data['action']: case 'list': # list tracks in album await socket.send(response.playlist( await ydl_wrap.get_playlist_items( ydls.get_ydl(data['site']), data['url'], ) )) case 'download': # download by URL ret = await ydl_wrap.download( ydls.get_ydl(data['site']), data['url'], data.get('items'), ) await socket.send(response.dl_end(ret)) # 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: cfg = config.get() async with websockets.serve(handler, cfg.host, cfg.port): await asyncio.Future() if __name__ == '__main__': asyncio.run(main(), debug=True)