musicdlp/backend/main.py

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)

            match data['action']:

                case 'list':  # list tracks in album
                    ydls.choose_ydl(data['site'])
                    await socket.send(response.playlist(
                        await ydls.get_playlist_items(data['url']),
                    ))

                case 'download':  # download by URL
                    ydls.choose_ydl(data['site'])
                    ret = await ydls.download(
                        data['url'],
                        data.get('items'),
                    )
                    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.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)