rollbot4/lib/rollbot/injection/args.py

from argparse import ArgumentParser, Namespace
from typing import Optional
import shlex

from ..types import Message, Context
from .base import Injector

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


class ArgsInjector(Injector[str]):
    async def inject(self, message: Message, context: Context) -> str:
        return message.command.args


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

    async def inject(self, message: Message, context: Context) -> str:
        if self.split is not None:
            return message.command.args.split(self.split)
        else:
            return message.command.args.split()


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

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


Args = ArgsInjector()
ArgList = ArgListSplitOn()
# TODO Arg(n)