|
24 | 24 | from collections.abc import Callable, Coroutine
|
25 | 25 | from dataclasses import dataclass, field
|
26 | 26 | from enum import Enum, auto
|
27 |
| -from typing import Any, ClassVar, NamedTuple, cast |
| 27 | +from typing import Any, ClassVar, NamedTuple, Protocol, cast, runtime_checkable |
28 | 28 |
|
29 | 29 | from discord.ext import tasks
|
30 | 30 | from loguru import logger
|
@@ -93,6 +93,13 @@ class CriticalTaskConfig:
|
93 | 93 | health_check_interval: float = 300.0 # 5 minutes
|
94 | 94 |
|
95 | 95 |
|
| 96 | +@runtime_checkable |
| 97 | +class CriticalTasksProvider(Protocol): |
| 98 | + """Protocol for cogs that provide critical tasks.""" |
| 99 | + |
| 100 | + def get_critical_tasks(self) -> list[CriticalTaskConfig]: ... |
| 101 | + |
| 102 | + |
96 | 103 | class TaskManager:
|
97 | 104 | """
|
98 | 105 | Enhanced task manager with health monitoring, metrics, and recovery capabilities.
|
@@ -205,20 +212,17 @@ def discover_and_register_cog_tasks(self) -> None:
|
205 | 212 | """
|
206 | 213 | Discover and register critical tasks from all loaded cogs.
|
207 | 214 |
|
208 |
| - This method asks each cog if it has critical tasks to register, |
209 |
| - making the system dynamic and cog-driven instead of hardcoded. |
| 215 | + This method iterates through all loaded cogs and looks for a |
| 216 | + `get_critical_tasks` method. If found, it calls the method to |
| 217 | + get a list of CriticalTaskConfig objects and registers them. |
210 | 218 | """
|
211 |
| - logger.info("Discovering critical tasks from cogs...") |
212 |
| - |
213 | 219 | for cog_name, cog in self.bot.cogs.items():
|
214 |
| - # Check if the cog has a method to report its critical tasks |
215 |
| - get_tasks_method = getattr(cog, "get_critical_tasks", None) |
216 |
| - if get_tasks_method and callable(get_tasks_method): |
| 220 | + if isinstance(cog, CriticalTasksProvider): |
217 | 221 | try:
|
218 |
| - if task_configs := get_tasks_method(): |
219 |
| - for config in task_configs: |
220 |
| - self.register_critical_task(config) |
221 |
| - logger.debug(f"Discovered task {config.name} from cog {cog_name}") |
| 222 | + task_configs = cog.get_critical_tasks() |
| 223 | + for config in task_configs: |
| 224 | + self.register_critical_task(config) |
| 225 | + logger.debug(f"Discovered task {config.name} from cog {cog_name}") |
222 | 226 | except Exception as e:
|
223 | 227 | logger.warning(f"Error discovering tasks from cog {cog_name}: {e}")
|
224 | 228 | continue
|
|
0 commit comments