space/plugin.py

from __future__ import annotations

import asyncio
import inspect
from collections.abc import Awaitable, Callable
from typing import Protocol

from pydantic import BaseModel

from .event import EventEnvelope

EventHandler = Callable[[BaseModel], Awaitable[None]]
QueuedEvent = tuple[EventEnvelope[BaseModel], EventHandler]


class PluginManagerProtocol(Protocol):
    def register_event(self, name: str, contract: type[BaseModel]) -> None: ...

    async def emit(
        self,
        name: str,
        data: BaseModel,
        *,
        source: str = "space",
    ) -> EventEnvelope[BaseModel]: ...

    async def wait_for_idle(self, *, timeout: float = 5) -> None: ...

    async def report_plugin_error(
        self,
        plugin: Plugin,
        envelope: EventEnvelope[BaseModel],
        error: Exception,
    ) -> None: ...


class Plugin:
    def __init__(self, name: str, *, receive_self_events: bool = False) -> None:
        if not name:
            raise ValueError("plugin name must not be empty")
        self.name: str = name
        self.capabilities: tuple[str, ...] = ()
        self.receive_self_events: bool = receive_self_events
        self.manager: PluginManagerProtocol | None = None
        self.enabled: bool = False
        self.queue: asyncio.Queue[QueuedEvent] = asyncio.Queue()
        self._initialized: bool = False
        self._worker: asyncio.Task[None] | None = None

    async def initialize(self) -> None:
        return None

    def iter_event_handlers(self) -> dict[str, EventHandler]:
        return {
            name.removeprefix("on_"): getattr(self, name)
            for name in dir(self)
            if name.startswith("on_") and inspect.iscoroutinefunction(getattr(self, name))
        }

    async def _enqueue(self, envelope: EventEnvelope[BaseModel], handler: EventHandler) -> None:
        if self.enabled:
            await self.queue.put((envelope, handler))

    def _start_worker(self) -> None:
        self._worker = asyncio.create_task(self._run(), name=f"space-plugin-{self.name}")

    async def _stop_worker(self) -> None:
        worker = self._worker
        self._worker = None
        if worker is not None:
            worker.cancel()
            try:
                await worker
            except asyncio.CancelledError:
                pass
        self.clear_queue()

    def clear_queue(self) -> None:
        while not self.queue.empty():
            self.queue.get_nowait()
            self.queue.task_done()

    async def _run(self) -> None:
        while True:
            envelope, handler = await self.queue.get()
            try:
                await handler(envelope.data)
            except Exception as error:
                if self.manager is not None:
                    await self.manager.report_plugin_error(self, envelope, error)
            finally:
                self.queue.task_done()