jarvis-core/jarvis_core/config.py
2023-03-23 22:06:12 -06:00

142 lines
4.3 KiB
Python

"""Load global config."""
import ast
from os import environ
from lib2to3.pgen2 import token
from pathlib import Path
from typing import Any, Dict, List, Optional, Type
import orjson as json
import yaml
from dotenv import load_dotenv
from pydantic import BaseModel
try:
from yaml import CLoader as Loader
except ImportError:
from yaml import Loader
DEFAULT_YAML = Path("config.yaml")
DEFAULT_JSON = Path("config.json")
class Config:
REQUIRED: Dict[str, Type] = {}
OPTIONAL: Dict[str, Dict[str, Any]] = {}
def __new__(cls, *args: list, **kwargs: dict):
"""Create a new Config object."""
for arg, flags in cls.OPTIONAL.items():
kwargs[arg] = kwargs.get(arg, flags["default"])
inst = super().__new__(cls, *args, **kwargs)
inst._validate()
return inst
def _validate(self) -> None:
for key in self.REQUIRED.keys():
if getattr(self, key, None) is None:
raise ValueError(f"Missing required key: {key}")
@classmethod
def from_env(cls) -> "Config":
"""Load the .env config file."""
load_dotenv()
data = {}
for k, t in cls.REQUIRED.items():
value = environ.get(k.upper(), None)
if value and not isinstance(value, t):
if t is dict:
try:
value = ast.literal_eval(value)
except Exception:
raise ValueError(f"{k} is a dict but is not formatted properly")
elif t is bool:
value = value.lower() in ["true", "yes", "t", "y", "1"]
else:
try:
value = t(value)
except Exception:
continue
data[k] = value
for k, flags in cls.OPTIONAL.items():
t = flags["type"]
value = environ.get(k.upper(), flags["default"])
if value and not isinstance(value, t):
if t is dict:
try:
value = ast.literal_eval(value)
except Exception:
raise ValueError(f"{k} is a dict but is not formatted properly")
elif t is bool:
value = value.lower() in ["true", "yes", "t", "y", "1"]
else:
try:
value = t(value)
except Exception:
continue
data[k] = value
return cls(**data)
@classmethod
def from_json(cls, filepath: Path | str = DEFAULT_JSON) -> "Config":
"""Load the json config file."""
if inst := cls.__dict__.get("inst"):
return inst
if isinstance(filepath, str):
filepath = Path(filepath)
with filepath.open() as f:
raw = f.read()
j = json.loads(raw)
return cls(**j)
@classmethod
def from_yaml(cls, filepath: Path | str = DEFAULT_YAML) -> "Config":
"""Load the yaml config file."""
if inst := cls.__dict__.get("inst"):
return inst
if isinstance(filepath, str):
filepath = Path(filepath)
with filepath.open() as f:
raw = f.read()
y = yaml.load(raw, Loader=Loader)
return cls(**y)
@classmethod
def load(cls, method: Optional[str] = None) -> "Config":
"""
Load the config in a somewhat generic way.
Default load order is: yaml, json, env
Args:
method: Load method, one of: yaml, json, env
"""
methods = ["yaml", "json", "env"]
if method and method in methods:
methods.remove(method)
methods.insert(0, method)
for method in methods:
try:
m = cls.__dict__.get(f"from_{method}", None)
if m:
return m()
except Exception:
continue
raise ValueError("Unable to load configuration, please create one of: config.yaml, config.json, .env")
@classmethod
def reload(cls) -> bool:
"""Reload the config."""
return cls.__dict__.pop("inst", None) is None