rollbot4/lib/rollbot/injection/util.py

from collections.abc import Callable, Coroutine
from typing import TypeVar, Any
from datetime import datetime

from aiohttp import ClientSession

from ..types import Message, Context, Attachment, Command
from .base import Injector, InjectorWithCleanup

__all__ = [
    "Lazy",
    "MessageInjector",
    "ContextInjector",
    "Origin",
    "Channel",
    "Sender",
    "Timestamp",
    "OriginAdmin",
    "ChannelAdmin",
    "Text",
    "Attachments",
    "CommandInjector",
    "Request",
    "Respond",
    "Config",
]


Dep = TypeVar("Dep")


class Simple(Injector[Dep]):
    def __init__(self, extract: Callable[[Message, Context], Dep]):
        self.extract = extract

    async def inject(self, message: Message, context: Context) -> Dep:
        return self.extract(message, context)


MessageInjector = Simple[Message](lambda m, c: m)
ContextInjector = Simple[Context](lambda m, c: c)
Origin = Simple[str](lambda m, c: m.origin_id)
Channel = Simple[str](lambda m, c: m.channel_id)
Sender = Simple[str](lambda m, c: m.sender_id)
Timestamp = Simple[datetime](lambda m, c: m.timestamp)
OriginAdmin = Simple[bool](lambda m, c: m.origin_admin)
ChannelAdmin = Simple[bool](lambda m, c: m.channel_admin)
Text = Simple[str](lambda m, c: m.text)
Attachments = Simple[list[Attachment]](lambda m, c: m.attachments)
CommandInjector = Simple[Command](lambda m, c: m.command)
Request = Simple[ClientSession](lambda m, c: c.request)
Respond = Simple[Callable[[], Coroutine[None, None, None]]](lambda m, c: c.respond)


class Config(Injector[Any]):
    def __init__(self, key: str):
        self.key = key

    async def inject(self, message: Message, context: Context) -> Any:
        return context.config(self.key)


class Lazy(InjectorWithCleanup[Callable[[], Coroutine[None, None, Dep]]]):
    def __init__(self, deferred: Injector[Dep]):
        self.deferred = deferred

    async def inject(
        self, message: Message, context: Context
    ) -> Callable[[], Coroutine[None, None, Dep]]:
        class _Wrapper:
            def __init__(self, deferred):
                self._calculated = None

                async def call():
                    if self._calculated is None:
                        self._calculated = await deferred.inject(message, context)
                    return self._calculated

                self._call = call

            def __call__(self):
                return self._call()

        return _Wrapper(self.deferred)

    async def cleanup(self, dep: Callable[[], Coroutine[None, None, Dep]]):
        if isinstance(self.deferred, InjectorWithCleanup) and dep._calculated is not None:
            await self.deferred.cleanup(dep._calculated)