py-dj/services.py

83 lines
2.7 KiB
Python
Raw Blame History

This file contains invisible Unicode characters

This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import asyncio
import yt_dlp
from fastapi import WebSocket
from pydantic import BaseModel
from models import Track, WSConnectionType
class ConnectionManager:
def __init__(self) -> None:
self.active_connections: dict[str, set[WebSocket]] = {
WSConnectionType.state.value: set(),
WSConnectionType.queue.value: set(),
}
async def connect(self, websocket: WebSocket, type: WSConnectionType):
await websocket.accept()
self.active_connections[type.value].add(websocket)
async def send(self, ws: WebSocket, message: BaseModel):
try:
await ws.send_json(message.model_dump())
except Exception:
self.disconnect(ws)
async def broadcast(self, ws_type: WSConnectionType, message: BaseModel):
broken = set()
conn_list = list(self.active_connections[ws_type.value])
for idx in range(len(conn_list)):
ws = conn_list[idx]
try:
await ws.send_json(message.model_dump())
except Exception:
broken.add(ws)
for ws in broken:
self.disconnect(ws)
def disconnect(self, websocket: WebSocket):
if websocket in self.active_connections[WSConnectionType.state.value]:
self.active_connections[WSConnectionType.state.value].remove(websocket)
if websocket in self.active_connections[WSConnectionType.queue.value]:
self.active_connections[WSConnectionType.queue.value].remove(websocket)
class DownloadService:
ydl_opts = {
"format": "mp3/bestaudio/best",
# See help(yt_dlp.postprocessor) for a list of available Postprocessors and their arguments
"postprocessors": [
{
"key": "FFmpegExtractAudio",
"preferredcodec": "mp3",
}
],
"paths": {"home": "queue"},
"outtmpl": {"default": "%(artist)s - %(track)s [%(id)s].%(ext)s"},
}
def __init__(self) -> None:
self.ydl = yt_dlp.YoutubeDL(self.ydl_opts)
async def download(self, url: str) -> Track:
def extract():
return self.ydl.extract_info(url, download=True)
info = await asyncio.to_thread(extract)
try:
filepath = info["requested_downloads"][-1]["filepath"] # type: ignore
except KeyError:
raise ValueError("Could not ")
track = Track(
artist=info.get("artist", None), # type: ignore
title=info.get("title", None), # type: ignore
duration=info.get("duration", None), # type: ignore
filepath=filepath, # type: ignore
)
print(f"Finished processing {track}")
return track