rollbot4/lib/rollbot/command/injection.py

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()