rollbot4/lib/rollbot/command/decorators.py

from collections.abc import Callable
from typing import Union
from functools import wraps
import inspect
import asyncio

from ..types import (
    Message,
    Context,
    CommandType,
    Response,
    StartupShutdownType,
    CommandConfiguration,
)
from .failure import RollbotFailureException
from .injection import Injector, inject_message, inject_context

decorated_startup: list[StartupShutdownType] = []
decorated_shutdown: list[StartupShutdownType] = []
decorated_commands: dict[str, CommandType] = {}


def on_startup(fn):
    decorated_startup.append(fn)
    return fn


def on_shutdown(fn):
    decorated_shutdown.append(fn)
    return fn


def as_command(arg: Union[str, Callable]):
    def impl(name, fn):
        if inspect.isasyncgenfunction(fn):
            lifted = fn
        elif inspect.iscoroutinefunction(fn):

            @wraps(fn)
            async def lifted(*args):
                yield await fn(*args)

        elif inspect.isgeneratorfunction(fn):

            @wraps(fn)
            async def lifted(*args):
                for res in fn(*args):
                    yield res

        elif inspect.isfunction(fn):

            @wraps(fn)
            async def lifted(*args):
                yield fn(*args)

        else:
            raise ValueError  # TODO details

        injectors = []
        for param in inspect.signature(fn).parameters:
            annot = fn.__annotations__[param]
            if annot == Message:
                injectors.append(inject_message)
            elif annot == Context:
                injectors.append(inject_context)
            elif isinstance(annot, Injector):
                injectors.append(annot.inject)
            else:
                raise ValueError  # TODO details

        async def command_impl(message: Message, context: Context):
            args = await asyncio.gather(*[inj(message, context) for inj in injectors])

            try:
                async for result in lifted(*args):
                    if isinstance(result, Response):
                        response = result
                    elif isinstance(result, str):
                        response = Response.from_message(message, text=result)
                    # TODO handle attachments, other special returns
                    else:
                        response = Response.from_message(message, str(result))
                    await context.respond(response)
            except RollbotFailureException as exc:
                # TODO handle errors more specifically
                await context.respond(Response.from_message(message, str(exc.failure)))

        decorated_commands[name] = command_impl
        return fn

    if isinstance(arg, str):
        return lambda fn: impl(arg, fn)
    else:
        return impl(arg.__name__, arg)


def get_command_config() -> CommandConfiguration:
    return CommandConfiguration(
        commands=decorated_commands,
        call_and_response={},
        aliases={},
        bangs=(),
        startup=decorated_startup,
        shutdown=decorated_shutdown,
    )