bot.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  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 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(
  26. *[task(self.context) for task in self.command_config.startup]
  27. )
  28. async def on_shutdown(self):
  29. await asyncio.gather(
  30. *[task(self.context) for task in self.command_config.shutdown]
  31. )
  32. async def on_message(self, incoming: RawMsg):
  33. message = self.parse(incoming)
  34. if message.text is None:
  35. return
  36. message.command = Command.from_text(message.text)
  37. if (
  38. message.command is None
  39. or message.command.bang not in self.command_config.bangs
  40. ):
  41. return
  42. command = self.command_config.aliases.get(
  43. message.command.name, message.command.name
  44. )
  45. res = self.command_config.call_and_response.get(command, None)
  46. if res is not None:
  47. await self.respond(Response.from_message(message, res))
  48. return
  49. command_call = self.command_config.commands.get(command, None)
  50. if command_call is None:
  51. await self.respond(
  52. Response.from_message(
  53. message, f"Sorry! I don't know the command {command}."
  54. )
  55. )
  56. return
  57. await command_call(message, self.context)