diff --git a/.flake8 b/.flake8 index 9417085..a5f0237 100644 --- a/.flake8 +++ b/.flake8 @@ -1,6 +1,7 @@ [flake8] exclude = tests/* + run.py extend-ignore = Q0, E501, C812, E203, W503, # These default to arguing with Black. We might configure some of them eventually diff --git a/jarvis_tasks/__init__.py b/jarvis_tasks/__init__.py index ccc538b..4bbf66a 100644 --- a/jarvis_tasks/__init__.py +++ b/jarvis_tasks/__init__.py @@ -1,2 +1,69 @@ """JARVIS background tasks.""" -__version__ = "0.1.0" +import asyncio +from typing import Optional + +from dis_snek import Intents, Snake +from jarvis_core.db import connect +from jarvis_core.log import get_logger + +from jarvis_tasks.config import TaskConfig +from jarvis_tasks.tasks import ban, reminder, twitter, warning + +logger = None +debug = False + + +async def _start(config: Optional[str] = "config.yaml") -> None: + """ + Main start function. + + Args: + config: Config path + """ + # Load config + config = TaskConfig.from_yaml(config) + + # Connect to database + testing = config.mongo["database"] != "jarvis" + logger.debug(f"Connecting to database, testing={testing}") + connect(**config.mongo["connect"], testing=testing) + + # Get event loop + loop = asyncio.get_event_loop() + + # Login as bot + logger.debug("Logging in bot") + intents = Intents.DEFAULT | Intents.GUILD_MEMBERS + bot = Snake(intents=intents, loop=loop) + await bot.login(config.token) + logger.info(f"Logged in as {bot.user.username}#{bot.user.discriminator}") + + # Start tasks + try: + logger.debug("Starting tasks") + functions = [ban.unban, reminder.remind, twitter.twitter, warning.unwarn] + tasks = [loop.create_task(f(bot, logger)) for f in functions] + for task in tasks: + await task + except KeyboardInterrupt: + for task in tasks: + task.cancel() + + +def start(config: Optional[str] = "config.yaml") -> None: + """ + Start the background tasks. + + Args: + config: Config path + """ + global logger, debug + # Set log level + _config = TaskConfig.from_yaml(config) + debug = _config.log_level == "DEBUG" + logger = get_logger(__name__, debug=debug) + logger.setLevel(_config.log_level) + + # Run the main tasks + logger.debug("Starting asyncio") + asyncio.run(_start(config)) diff --git a/poetry.lock b/poetry.lock index f038eab..05dd97c 100644 --- a/poetry.lock +++ b/poetry.lock @@ -200,7 +200,7 @@ plugins = ["setuptools"] [[package]] name = "jarvis-core" -version = "0.3.0" +version = "0.4.1" description = "" category = "main" optional = false @@ -218,7 +218,7 @@ umongo = "^3.1.0" type = "git" url = "https://git.zevaryx.com/stark-industries/jarvis/jarvis-core.git" reference = "main" -resolved_reference = "b8db789590bc2b5148e19f15e86167403b3e8cc3" +resolved_reference = "f57d79c01a1781bf299b7ea504db60767a808545" [[package]] name = "jedi" diff --git a/run.py b/run.py new file mode 100644 index 0000000..6ced8f7 --- /dev/null +++ b/run.py @@ -0,0 +1,4 @@ +from jarvis_tasks import start + +if __name__ == "__main__": + start()