src/tasks/__init__.py

"""Task registry. To register a new task:
    1. Create a subpackage under src/tasks/<your_task>/ exposing TASK (see _template/).
    2. Import it below and add to TASKS.
"""
from typing import Dict

from src.tasks.base import TaskPlugin
from src.tasks.task_future_prediction import TASK as TASK_FUTURE_PRED
from src.tasks.task_property_estimation import TASK as TASK_PROPERTY_EST

TASKS: Dict[str, TaskPlugin] = {
    TASK_FUTURE_PRED.name: TASK_FUTURE_PRED,
    TASK_PROPERTY_EST.name: TASK_PROPERTY_EST,
}


def get_task(name: str) -> TaskPlugin:
    if name not in TASKS:
        raise ValueError(f"Unknown task: {name}")
    return TASKS[name]