KuroCore/kurocore/bot.py

import sys
from asyncio import set_event_loop_policy
from platform import system

from .logger import BotLogger, RequestLogger, ClientLogger
from .main.handler import Handler
from kurocore.database import Database


class Bot:
    __slots__ = ('config', 'is_running', 'middlewares', 'tasks')

    def __init__(self, config, load_db=False):
        BotLogger(config)
        RequestLogger(config)
        ClientLogger(config)
        self.config = config
        if load_db:
            Database(self.config)
        self.is_running = False
        self.middlewares = []
        self.tasks = []
        if sys.version_info >= (3, 8) and sys.platform.lower().startswith("win"):
            from asyncio import WindowsSelectorEventLoopPolicy
            set_event_loop_policy(WindowsSelectorEventLoopPolicy())

    def add_plugin(self, plugin):
        if self.is_running:
            BotLogger.log.error('Bot already running!')
        self.config.plugins.append(plugin)

    def add_prefix(self, prefix):
        if self.is_running:
            BotLogger.log.error('Bot already running!')
        self.config.prefixes.append(prefix)

    def add_middleware(self, middleware):
        if self.is_running:
            BotLogger.log.error('Bot already running!')
        self.middlewares.append(middleware)

    def add_task(self, task):
        if self.is_running:
            BotLogger.log.error('Bot already running!')
        if system == 'Windows':
            BotLogger.log.warning('Windows don\'t support tasks')
            return
        self.tasks.append(task)

    def run(self):
        handler = Handler(self.config, self.middlewares)
        handler.init()
        try:
            for task in self.tasks:
                BotLogger.log.debug(f'registered task: {task}')
                handler.loop.create_task(task())
            self.is_running = True
            handler.run()
        except KeyboardInterrupt:
            handler.shutdown()
            if Database.db and not Database.db.is_closed():
                Database.db.close()