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
  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 NotImplemented("Must be implemented by driver")
  20. def parse(self, incoming: RawMsg) -> Message:
  21. raise NotImplemented("Must be implemented by driver")
  22. async def respond(self, response: Response):
  23. raise NotImplemented("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. cleaned = message.text.lstrip()
  33. if len(cleaned) == 0 or cleaned[0] not in self.command_config.bangs:
  34. return
  35. parts = cleaned[1:].lstrip().split(maxsplit=1)
  36. if len(parts) == 0:
  37. return
  38. command = self.command_config.aliases.get(parts[0], parts[0])
  39. res = self.command_config.call_and_response.get(command, None)
  40. if res is not None:
  41. await self.respond(Response.from_message(message, res))
  42. return
  43. command_call = self.command_config.commands.get(command, None)
  44. if command is None:
  45. await self.respond(Response.from_message(message, f"Sorry! I don't know the command {command}."))
  46. return
  47. await command_call(message, self.context)