from typing import Generic, TypeVar, Optional from argparse import ArgumentParser, Namespace from collections.abc import Callable, Coroutine import shlex from aiosqlite.core import Connection from ..types import Message, Context async def inject_message(message: Message, context: Context) -> Message: return message async def inject_context(message: Message, context: Context) -> Context: return context Dep = TypeVar("DepType") class Injector(Generic[Dep]): async def inject(self, message: Message, context: Context) -> Dep: raise NotImplemented 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)) class DatabaseInjector(Injector[Connection]): async def inject(self, message: Message, context: Context) -> Connection: return context.database() class Lazy(Injector[Callable[[], Coroutine[None, None, Dep]]]): def __init__(self, deferred: Injector[Dep]): self.deferred = deferred self._calculated = None async def inject(self, message: Message, context: Context) -> Callable[[], Coroutine[None, None, Dep]]: async def call(): if self._calculated is None: self._calculated = await self.deferred.inject(message, context) return self._calculated return call Args = ArgsInjector() ArgList = ArgListSplitOn() Database = DatabaseInjector() # TODO data query handler