musicdlp/backend/main.py

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)