bot.py 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. from dataclasses import dataclass
  2. from typing import Any, Generic, TypeVar
  3. import asyncio
  4. import aiosqlite
  5. from .types import CommandConfiguration, Message, Response, Context, Command
  6. # TODO logging
  7. RawMsg = TypeVar("RawMsg")
  8. @dataclass
  9. class Rollbot(Generic[RawMsg]):
  10. command_config: CommandConfiguration
  11. database_file: str
  12. def __post_init__(self):
  13. self.context = Context(
  14. config=self.read_config,
  15. respond=self.respond,
  16. database=lambda: aiosqlite.connect(self.database_file),
  17. )
  18. def read_config(self, key: str) -> Any:
  19. raise NotImplementedError("Must be implemented by driver")
  20. def parse(self, incoming: RawMsg) -> Message:
  21. raise NotImplementedError("Must be implemented by driver")
  22. async def respond(self, response: Response):
  23. raise NotImplementedError("Must be implemented by driver")
  24. async def on_startup(self):
  25. await asyncio.gather(*[task(self.context) for task in self.command_config.startup])
  26. async def on_shutdown(self):
  27. await asyncio.gather(*[task(self.context) for task in self.command_config.shutdown])
  28. async def on_message(self, incoming: RawMsg):
  29. message = self.parse(incoming)
  30. if message.text is None:
  31. return
  32. message.command = Command.from_text(message.text)
  33. if message.command is None or message.command.bang not in self.command_config.bangs:
  34. return
  35. command = self.command_config.aliases.get(message.command.name, message.command.name)
  36. res = self.command_config.call_and_response.get(command, None)
  37. if res is not None:
  38. await self.respond(Response.from_message(message, res))
  39. return
  40. command_call = self.command_config.commands.get(command, None)
  41. if command_call is None:
  42. await self.respond(
  43. Response.from_message(message, f"Sorry! I don't know the command {command}.")
  44. )
  45. return
  46. await command_call(message, self.context)