from typing_extensions import TypeVar

from .base_defense import Defense

DEFENSES: dict[str, type[Defense]] = {}

T = TypeVar("T", bound=Defense)


def register_defense(defense: type[T]) -> type[T]:
    """Register a defense subclass with the defense_hparams module."""
    DEFENSES[defense.__name__] = defense
    return defense


def load_defense(name: str) -> type[Defense]:
    """Load a defense subclass by name."""
    return DEFENSES[name]
