bot.py 2.1 KB

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