rollbot4/lib/rollbot/injection/util.py

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

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

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


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(lambda m, c: m)
ContextInjector = Simple(lambda m, c: c)
Origin = Simple(lambda m, c: m.origin_id)
Channel = Simple(lambda m, c: m.channel_id)
Sender = Simple(lambda m, c: m.sender_id)
Timestamp = Simple(lambda m, c: m.timestamp)
OriginAdmin = Simple(lambda m, c: m.origin_admin)
ChannelAdmin = Simple(lambda m, c: m.channel_admin)
Text = Simple(lambda m, c: m.text)
Attachments = Simple(lambda m, c: m.attachments)
CommandInjector = Simple(lambda m, c: m.command)


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)