rollbot4/rollbot/rollbot/types.py

from __future__ import annotations

from logging import Logger
from dataclasses import dataclass, field
from datetime import datetime
from collections.abc import Callable, Coroutine, Container
from typing import Union, Any, Optional
import time

from aiosqlite import Connection
from aiohttp import ClientSession

__all__ = [
    "Attachment",
    "Message",
    "Command",
    "Response",
    "Context",
    "CommandType",
    "StartupShutdownType",
    "CommandConfiguration",
]


@dataclass
class Attachment:
    name: str
    body: Union[str, bytes, None, Callable[[], Coroutine[Any, None, Any]]]


@dataclass
class Message:
    origin_id: str
    channel_id: str
    sender_id: str
    timestamp: datetime
    origin_admin: bool
    channel_admin: bool
    sender_name: Optional[str] = None
    text: Optional[str] = None
    attachments: list[Attachment] = field(default_factory=list)
    message_id: Optional[str] = None
    received_at: float = field(default_factory=time.time)

    def __post_init__(self):
        self.command = None


@dataclass
class Command:
    bang: str
    name: str
    raw_name: str
    args: str

    def __post_init__(self):
        # cache is used by injectors to store results and avoid recomputation
        self.cache = {}

    @staticmethod
    def from_text(text: str) -> Optional[Command]:
        cleaned = text.lstrip()

        if len(cleaned) < 2:
            return None

        parts = cleaned[1:].lstrip().split(maxsplit=1)
        if len(parts) == 0:
            return None

        return Command(
            bang=cleaned[0],
            name=parts[0].lower(),
            raw_name=parts[0],
            args=parts[1] if len(parts) > 1 else "",
        )

    def get_subcommand(self, inherit_bang=True) -> Command:
        saved = self.cache.get(("subcommand", inherit_bang), None)
        if saved is None:
            if inherit_bang and not self.args.startswith(self.bang):
                saved = Command.from_text(self.bang + self.args)
            else:
                saved = Command.from_text(self.args)
            self.cache[("subcommand", inherit_bang)] = saved
        return saved


@dataclass
class Response:
    origin_id: str
    channel_id: str
    text: Optional[str] = None
    attachments: Optional[list[Attachment]] = None
    cause: Optional[Message] = None

    @staticmethod
    def from_message(
        msg: Message, text: Optional[str] = None, attachments: list[Attachment] = None
    ) -> Response:
        return Response(
            origin_id=msg.origin_id,
            channel_id=msg.channel_id,
            text=text,
            attachments=attachments or [],
            cause=msg,
        )


@dataclass
class Context:
    config: Callable[[str], Any]
    respond: Callable[[], Coroutine[None, None, None]]
    request: ClientSession
    database: Callable[[], Coroutine[None, None, Connection]]
    logger: Logger
    debugging: Optional[str] = None

    def get_debugging(self) -> Optional[str]:
        old = self.debugging
        self.debugging = None
        return old


CommandType = Callable[[Message, Context], Coroutine[None, None, None]]
StartupShutdownType = Callable[[Context], Coroutine[None, None, None]]


@dataclass
class CommandConfiguration:
    commands: dict[str, CommandType] = field(default_factory=dict)
    call_and_response: dict[str, str] = field(default_factory=dict)
    aliases: dict[str, str] = field(default_factory=dict)
    bangs: Container[str] = ("!",)
    startup: list[StartupShutdownType] = field(default_factory=list)
    shutdown: list[StartupShutdownType] = field(default_factory=list)

    def extend(self, other: CommandConfiguration) -> CommandConfiguration:
        return CommandConfiguration(
            commands={
                **self.commands,
                **other.commands,
            },
            call_and_response={
                **self.call_and_response,
                **other.call_and_response,
            },
            aliases={
                **self.aliases,
                **other.aliases,
            },
            bangs=(*self.bangs, *other.bangs),
            startup=[*self.startup, *other.startup],
            shutdown=[*self.shutdown, *other.shutdown],
        )