import asyncio
from typing import Coroutine


class TaskManager:
    """ 任务管理器 """
    def __init__(self):
        self.tasks: list[asyncio.Task] = []
        
    def create_task(self, coro: Coroutine):
        """ 创建任务 """
        task = asyncio.create_task(coro)
        task.add_done_callback(self.tasks.remove)
        self.tasks.append(task)
        return task
    
    def cancel_all(self):
        """ 取消所有任务 """
        for task in self.tasks:
            task.cancel()
        self.tasks.clear()
        
    async def wait_all(self):
        """ 等待所有任务完成 """
        await asyncio.gather(*self.tasks)
        self.tasks.clear()
        