mirror of
https://github.com/wahyd4/you-music.git
synced 2026-08-08 20:59:47 +10:00
83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
from datetime import datetime
|
|
from typing import Dict, List, Optional
|
|
from pydantic import BaseModel
|
|
import asyncio
|
|
|
|
|
|
class DownloadTask(BaseModel):
|
|
id: str
|
|
url: str
|
|
title: str
|
|
status: str # pending, downloading, completed, failed
|
|
progress: float = 0.0
|
|
error: Optional[str] = None
|
|
music_id: Optional[int] = None
|
|
created_at: datetime = datetime.now()
|
|
updated_at: datetime = datetime.now()
|
|
thumbnail: Optional[str] = None
|
|
artist: Optional[str] = None
|
|
|
|
class Config:
|
|
arbitrary_types_allowed = True
|
|
|
|
|
|
class DownloadQueue:
|
|
def __init__(self):
|
|
self.tasks: Dict[str, DownloadTask] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def add_task(self, task_id: str, url: str, title: str, thumbnail: Optional[str] = None, artist: Optional[str] = None) -> DownloadTask:
|
|
async with self._lock:
|
|
task = DownloadTask(
|
|
id=task_id,
|
|
url=url,
|
|
title=title,
|
|
status="pending",
|
|
thumbnail=thumbnail,
|
|
artist=artist
|
|
)
|
|
self.tasks[task_id] = task
|
|
return task
|
|
|
|
async def update_status(self, task_id: str, status: str, progress: float = None, error: str = None, music_id: int = None):
|
|
async with self._lock:
|
|
if task_id in self.tasks:
|
|
task = self.tasks[task_id]
|
|
task.status = status
|
|
task.updated_at = datetime.now()
|
|
if progress is not None:
|
|
task.progress = progress
|
|
if error is not None:
|
|
task.error = error
|
|
if music_id is not None:
|
|
task.music_id = music_id
|
|
|
|
async def get_task(self, task_id: str) -> Optional[DownloadTask]:
|
|
async with self._lock:
|
|
return self.tasks.get(task_id)
|
|
|
|
async def get_all_tasks(self) -> List[DownloadTask]:
|
|
async with self._lock:
|
|
return list(self.tasks.values())
|
|
|
|
async def get_tasks_by_status(self, status: str) -> List[DownloadTask]:
|
|
async with self._lock:
|
|
return [task for task in self.tasks.values() if task.status == status]
|
|
|
|
async def remove_task(self, task_id: str):
|
|
async with self._lock:
|
|
if task_id in self.tasks:
|
|
del self.tasks[task_id]
|
|
|
|
async def clear_completed(self):
|
|
async with self._lock:
|
|
self.tasks = {
|
|
task_id: task
|
|
for task_id, task in self.tasks.items()
|
|
if task.status not in ["completed", "failed"]
|
|
}
|
|
|
|
|
|
# Global instance
|
|
download_queue = DownloadQueue()
|