rollbot4/lib/rollbot/injection/subcommand.py

from collections.abc import Callable
from typing import TypeVar, Optional

from ..types import Message, Context, Command
from ..failure import RollbotFailure
from .base import Injector, Simple
from .args import ArgListSplitOn, ArgParse, Arg

__all__ = [
    "Subcommand",
]

T = TypeVar("T")


def get_subc_args(msg: Message) -> str:
    if (sc := msg.command.get_subcommand()) is not None:
        return sc.args
    RollbotFailure.MISSING_SUBCOMMAND.raise_exc()


class SubcArgListSplitOn(ArgListSplitOn):
    def get_args(self, message: Message) -> str:
        return get_subc_args(message)


class SubcArgParse(ArgParse):
    def get_args(self, message: Message) -> str:
        return get_subc_args(message)


SubcArgList = SubcArgListSplitOn()


class SubcArg(Arg):
    def arg_source(self) -> Injector[list[str]]:
        return SubcArgList


class _Subcommand(Simple[Optional[Command]]):
    Args = Simple[Optional[str]](lambda m, c: get_subc_args(m))
    ArgList = SubcArgList
    ArgListSplitOn = SubcArgListSplitOn
    ArgParse = SubcArgParse
    Arg = SubcArg

    def __init__(self):
        super().__init__(lambda m, c: m.command.get_subcommand())


Subcommand = _Subcommand()