rollbot4/lib/rollbot/types.py

from __future__ import annotations

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

from aiosqlite.core import Connection

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


@dataclass
class Attachment:
    name: str
    body: Union[str, bytes]


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

    def __post_init__(self):
        self.command = None


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

    @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],
            args=parts[1] if len(parts) > 1 else "",
        )


@dataclass
class Response:
    origin_id: str
    channel_id: str
    text: Optional[str] = None
    attachments: Optional[list[Attachment]] = 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 [],
        )


@dataclass
class Context:
    config: Callable[[str], Any]
    respond: Callable[[], Coroutine[None, None, None]]
    database: Callable[[], Coroutine[None, None, Connection]]


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