rollbot4/lib/rollbot/injection/args.py

from argparse import ArgumentParser, Namespace
from typing import Optional, TypeVar
from collections.abc import Callable
import shlex

from ..types import Message, Context
from ..failure import RollbotFailure
from .base import Injector, Simple

__all__ = [
    "Args",
    "ArgList",
    "ArgListSplitOn",
    "ArgParse",
    "Arg",
]


Args = Simple[str](lambda m, c: m.command.args)


class ArgAccessorBase:
    def get_args(self, message: Message) -> str:
        return message.command.args


class ArgListSplitOn(Injector[list[str]], ArgAccessorBase):
    def __init__(self, split: Optional[str] = None):
        self.split = split

    async def inject(self, message: Message, context: Context) -> str:
        cache_key = (ArgListSplitOn.__name__, self.split)
        result = message.command.cache.get(cache_key, None)
        if result is not None:
            return result
        if self.split is not None:
            result = self.get_args(message).split(self.split)
        else:
            result = self.get_args(message).split()
        message.command.cache[cache_key] = result
        return result


class ArgParse(Injector[Namespace], ArgAccessorBase):
    def __init__(self, parser: ArgumentParser):
        self.parser = parser

    async def inject(self, message: Message, context: Context) -> Namespace:
        return self.parser.parse_args(shlex.split(self.get_args(message)))


ArgList = ArgListSplitOn()

ArgType = TypeVar("ArgType")


class Arg(Injector[ArgType]):
    def __init__(
        self,
        index: int = 0,
        convert: Callable[[str], ArgType] = str,
        required: bool = True,
        default: Optional[ArgType] = None,
        missing_msg: Optional[str] = None,
        fail_msg: Optional[str] = None,
    ):
        self.index = index
        self.convert = convert
        self.required = required
        self.default = default
        self.missing_msg = missing_msg or f"Missing argument {self.index}"
        self.fail_msg = fail_msg or "Invalid argument: {}"

    def arg_source(self) -> Injector[list[str]]:
        return ArgList

    async def inject(self, message: Message, context: Context) -> str:
        try:
            arg = (await self.arg_source().inject(message, context))[self.index]
        except IndexError:
            if self.required:
                RollbotFailure.INVALID_ARGUMENTS.raise_exc(detail=self.missing_msg)
            else:
                return self.default
        try:
            return self.convert(arg)
        except ValueError:
            RollbotFailure.INVALID_ARGUMENTS.raise_exc(detail=self.fail_msg.format(arg))