bot.py 2.2 KB

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