argos/argos/server/settings.py

import os
from functools import lru_cache
from os import environ

import yaml
from pydantic_settings import BaseSettings, SettingsConfigDict
from yamlinclude import YamlIncludeConstructor

from argos.schemas.config import Config


class Settings(BaseSettings):
    model_config = SettingsConfigDict(env_prefix="argos_", env_file=".env")
    app_env: str
    database_url: str
    yaml_file: str


class DevSettings(Settings):
    app_env: str = "dev"
    database_url: str = "sqlite:////tmp/argos.db"
    yaml_file: str = "config.yaml"


class TestSettings(Settings):
    app_env: str = "test"
    yaml_file: str = "tests/config.yaml"
    database_url: str = "sqlite:////tmp/test-argos.db"


class ProdSettings(Settings):
    app_env: str = "prod"
    pass


environments = {
    "dev": DevSettings,
    "prod": ProdSettings,
    "test": TestSettings,
}


@lru_cache()
def get_app_settings() -> Settings:
    app_env = environ.get("ARGOS_APP_ENV", "dev")
    settings = environments.get(app_env)
    return settings()


def read_yaml_config(filename):
    parsed = _load_yaml(filename)
    return Config(**parsed)


def _load_yaml(filename):
    base_dir = os.path.dirname(filename)
    YamlIncludeConstructor.add_to_loader_class(
        loader_class=yaml.FullLoader, base_dir=base_dir
    )

    with open(filename, "r") as stream:
        return yaml.load(stream, Loader=yaml.FullLoader)