107 lines
2.5 KiB
Python
107 lines
2.5 KiB
Python
"""Task config."""
|
|
from enum import Enum
|
|
from os import environ
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import orjson as json
|
|
import yaml
|
|
from dotenv import load_dotenv
|
|
from jarvis_core.util import find_all
|
|
from pydantic import BaseModel
|
|
|
|
try:
|
|
from yaml import CLoader as Loader
|
|
except ImportError:
|
|
from yaml import Loader
|
|
|
|
|
|
class Environment(Enum):
|
|
"""JARVIS running environment."""
|
|
|
|
production = "production"
|
|
develop = "develop"
|
|
|
|
|
|
class Mongo(BaseModel):
|
|
"""MongoDB config."""
|
|
|
|
host: list[str] | str = "localhost"
|
|
username: Optional[str] = None
|
|
password: Optional[str] = None
|
|
port: int = 27017
|
|
|
|
|
|
class Config(BaseModel):
|
|
"""Tasks config model."""
|
|
|
|
token: str
|
|
mongo: Mongo
|
|
log_level: str = "INFO"
|
|
environment: Environment = Environment.develop
|
|
|
|
|
|
_config: Config = None
|
|
|
|
|
|
def _load_json() -> Config | None:
|
|
path = Path("config.json")
|
|
if path.exists():
|
|
with path.open() as f:
|
|
j = json.loads(f.read())
|
|
return Config(**j)
|
|
|
|
|
|
def _load_yaml() -> Config | None:
|
|
path = Path("config.yaml")
|
|
if path.exists():
|
|
with path.open() as f:
|
|
y = yaml.load(f.read(), Loader=Loader)
|
|
return Config(**y)
|
|
|
|
|
|
def _load_env() -> Config | None:
|
|
load_dotenv()
|
|
data = {}
|
|
mongo = {}
|
|
mongo_keys = find_all(lambda x: x.upper().startswith("MONGO"), environ.keys())
|
|
|
|
config_keys = mongo_keys + ["TOKEN", "LOG_LEVEL", "ENVIRONMENT"]
|
|
|
|
for item, value in environ.items():
|
|
if item not in config_keys:
|
|
continue
|
|
|
|
if item in mongo_keys:
|
|
key = "_".join(item.split("_")[1:]).lower()
|
|
mongo[key] = value
|
|
else:
|
|
data[item.lower()] = value
|
|
|
|
data["mongo"] = mongo
|
|
|
|
return Config(**data)
|
|
|
|
|
|
def load_config(method: Optional[str] = None) -> Config:
|
|
"""
|
|
Load the config using the specified method first
|
|
|
|
Args:
|
|
method: Method to use first
|
|
"""
|
|
global _config
|
|
if _config is not None:
|
|
return _config
|
|
|
|
methods = {"yaml": _load_yaml, "json": _load_json, "env": _load_env}
|
|
method_names = list(methods.keys())
|
|
if method and method in method_names:
|
|
method_names.remove(method)
|
|
method_names.insert(0, method)
|
|
|
|
for method in method_names:
|
|
if _config := methods[method]():
|
|
return _config
|
|
|
|
raise FileNotFoundError("Missing one of: config.yaml, config.json, .env")
|