rollbot4/rollbot/rollbot/bot.py

from dataclasses import dataclass
from typing import Any, Generic, TypeVar
import traceback
import asyncio
import logging

import aiosqlite
import aiohttp

from .types import CommandConfiguration, Message, Response, Context, Command


RawMsg = TypeVar("RawMsg")


@dataclass
class Rollbot(Generic[RawMsg]):
    command_config: CommandConfiguration
    database_file: str

    def __post_init__(self):
        self.context = Context(
            config=self.read_config,
            respond=self.respond,
            request=aiohttp.ClientSession(),
            database=lambda: aiosqlite.connect(self.database_file),
            logger=logging.getLogger(),
        )

    def read_config(self, key: str) -> Any:
        raise NotImplementedError("Must be implemented by driver")

    def parse(self, incoming: RawMsg) -> Message:
        raise NotImplementedError("Must be implemented by driver")

    async def respond(self, response: Response):
        raise NotImplementedError("Must be implemented by driver")

    async def on_startup(self):
        await asyncio.gather(*[task(self.context) for task in self.command_config.startup])

    async def on_shutdown(self):
        await asyncio.gather(*[task(self.context) for task in self.command_config.shutdown])
        await self.context.request.close()

    async def on_message(self, incoming: RawMsg):
        message = self.parse(incoming)
        if message.text is None:
            return

        message.command = Command.from_text(message.text)
        if message.command is None or message.command.bang not in self.command_config.bangs:
            return

        command = self.command_config.aliases.get(message.command.name, message.command.name)

        if command == "help":
            args = message.command.args.split(maxsplit=1)
            if len(args) == 0:
                help_msg = "Use this command to look up the help info of another command, try !help roll for example!"
            else:
                cmd = self.command_config.commands.get(args[0], None)
                if cmd is None or cmd.__doc__ is None:
                    help_msg = "Sorry! I don't have any explanation of that command"
                else:
                    help_msg = cmd.__doc__.strip()
            await self.respond(Response.from_message(message, help_msg))
            return

        res = self.command_config.call_and_response.get(command, None)
        if res is not None:
            await self.respond(Response.from_message(message, res))
            return

        command_call = self.command_config.commands.get(command, None)
        if command_call is None:
            await self.respond(
                Response.from_message(message, f"Sorry! I don't know the command {command}.")
            )
            return

        try:
            await command_call(message, self.context)
        except Exception as e:
            self.context.logger.exception(
                f"Unhandled error during execution of {command} in {message}"
            )
            await self.respond(
                Response.from_message(
                    message, f"Encountered an unhandled error: {type(e).__name__}"
                )
            )
            self.context.debugging = "".join(traceback.format_exc())