rollbot4/lib/rollbot/command/decorators.py

from collections.abc import Callable
from typing import Union
import inspect

from ..types import (
    Message,
    Context,
    CommandType,
    Response,
    StartupShutdownType,
    CommandConfiguration,
)
from .failure import RollbotFailureException

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

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

        elif inspect.isgeneratorfunction(fn):

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

        elif inspect.isfunction(fn):

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

        else:
            raise ValueError  # TODO details

        async def command_impl(message: Message, context: Context):
            args = []  # TODO implement dep injection

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