import asyncio import json import websockets import response import ydl_pool import ydl_wrap type SocketT = websockets.WebSocketServerProtocol # TODO: config HOST = '127.0.0.1' PORT = 5678 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.ok_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.ok_downloaded(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: async with websockets.serve(handler, HOST, PORT): await asyncio.Future() if __name__ == '__main__': asyncio.run(main(), debug=True)