source/bimanual/yam/registry.py

"""Task registry, in the shape ManiSkill uses: one class per task, registered by name.

    @register_task("grape_box")
    class GrapeBoxTask(YamTask):
        ...

`make("grape_box")` then builds it. Importing `bimanual.yam_tasks.tasks` pulls in every task
module, which is what fills this registry -- the same trick as ManiSkill's `envs/tasks/__init__`.
"""
from __future__ import annotations

REGISTRY: dict[str, type] = {}


def register_task(name: str, **meta):
    """Class decorator: put a task class in the registry under `name`."""
    def deco(cls):
        if name in REGISTRY:
            raise ValueError(f"task {name!r} is already registered by {REGISTRY[name].__name__}")
        cls.task_name = name
        cls.meta = meta
        REGISTRY[name] = cls
        return cls
    return deco


def make(name: str, **kwargs):
    if name not in REGISTRY:
        raise SystemExit(f"unknown task {name!r}. Registered: {sorted(REGISTRY)}")
    return REGISTRY[name](**kwargs)


def list_tasks():
    rows = []
    for n, cls in sorted(REGISTRY.items()):
        rows.append({"name": n, "class": cls.__name__,
                     "title": getattr(cls, "title", ""),
                     "tags": ",".join(getattr(cls, "tags", [])),
                     "skill": getattr(cls, "skill", "")})
    return rows