102 lines
2.8 KiB
Python
102 lines
2.8 KiB
Python
"""Load the config for JARVIS"""
|
|
import os
|
|
|
|
from jarvis_core.config import Config as CConfig
|
|
from pymongo import MongoClient
|
|
from yaml import load
|
|
|
|
try:
|
|
from yaml import CLoader as Loader
|
|
except ImportError:
|
|
from yaml import Loader
|
|
|
|
|
|
class JarvisConfig(CConfig):
|
|
REQUIRED = ["token", "mongo", "urls"]
|
|
OPTIONAL = {
|
|
"sync": False,
|
|
"log_level": "WARNING",
|
|
"scales": None,
|
|
"events": True,
|
|
"gitlab_token": None,
|
|
"max_messages": 1000,
|
|
"twitter": None,
|
|
"reddit": None,
|
|
"rook_token": None,
|
|
}
|
|
|
|
|
|
class Config(object):
|
|
"""Config singleton object for JARVIS"""
|
|
|
|
def __new__(cls, *args: list, **kwargs: dict):
|
|
"""Get the singleton config, or creates a new one."""
|
|
it = cls.__dict__.get("it")
|
|
if it is not None:
|
|
return it
|
|
cls.__it__ = it = object.__new__(cls)
|
|
it.init(*args, **kwargs)
|
|
return it
|
|
|
|
def init(
|
|
self,
|
|
token: str,
|
|
mongo: dict,
|
|
urls: dict,
|
|
sync: bool = False,
|
|
log_level: str = "WARNING",
|
|
cogs: list = None,
|
|
events: bool = True,
|
|
gitlab_token: str = None,
|
|
max_messages: int = 1000,
|
|
twitter: dict = None,
|
|
reddit: dict = None,
|
|
rook_token: str = None,
|
|
) -> None:
|
|
"""Initialize the config object."""
|
|
self.token = token
|
|
self.mongo = mongo
|
|
self.urls = urls
|
|
self.log_level = log_level
|
|
self.cogs = cogs
|
|
self.events = events
|
|
self.max_messages = max_messages
|
|
self.gitlab_token = gitlab_token
|
|
self.twitter = twitter
|
|
self.reddit = reddit
|
|
self.sync = sync or os.environ.get("SYNC_COMMANDS", False)
|
|
self.rook_token = rook_token
|
|
self.__db_loaded = False
|
|
self.__mongo = MongoClient(**self.mongo["connect"])
|
|
|
|
def get_db_config(self) -> None:
|
|
"""Load the database config objects."""
|
|
if not self.__db_loaded:
|
|
db = self.__mongo[self.mongo["database"]]
|
|
items = db.config.find()
|
|
for item in items:
|
|
setattr(self, item["key"], item["value"])
|
|
self.__db_loaded = True
|
|
|
|
@classmethod
|
|
def from_yaml(cls, y: dict) -> "Config":
|
|
"""Load the yaml config file."""
|
|
return cls(**y)
|
|
|
|
|
|
def get_config(path: str = "config.yaml") -> Config:
|
|
"""Get the config from the specified yaml file."""
|
|
if Config.__dict__.get("it"):
|
|
return Config()
|
|
with open(path) as f:
|
|
raw = f.read()
|
|
y = load(raw, Loader=Loader)
|
|
config = Config.from_yaml(y)
|
|
config.get_db_config()
|
|
return config
|
|
|
|
|
|
def reload_config() -> None:
|
|
"""Force reload of the config singleton on next call."""
|
|
if "it" in Config.__dict__:
|
|
Config.__dict__.pop("it")
|