Spaces:
Running on Zero
Running on Zero
File size: 2,889 Bytes
35a216d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | 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()
|