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)