rollbot4/lib/rollbot/injection/util.py

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

from aiohttp import ClientSession

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

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

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)
# TODO might be nice to auto format with command name
LoggerInjector = Simple[Logger](lambda m, c: c.logger)


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)


Dep = TypeVar("Dep")


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)