jarvis-tasks/jarvis_tasks/__init__.py

77 lines
2 KiB
Python

"""JARVIS background tasks."""
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 import const
from jarvis_tasks.config import TaskConfig
from jarvis_tasks.tasks import ban, lock, lockdown, reddit, reminder, twitter, warning
__version__ = const.__version__
logger = None
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,
lock.unlock,
lockdown.lift,
reddit.reddit,
reminder.remind,
twitter.twitter,
warning.unwarn,
]
tasks = [loop.create_task(f(bot)) 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)
logger = get_logger(__name__)
logger.setLevel(_config.log_level)
# Run the main tasks
logger.debug("Starting asyncio")
asyncio.run(_start(config))