diff options
| -rw-r--r-- | bot/constants.py | 71 | ||||
| -rw-r--r-- | config-default.yml | 7 |
2 files changed, 42 insertions, 36 deletions
diff --git a/bot/constants.py b/bot/constants.py index 43f03d7bf..2a9796cc5 100644 --- a/bot/constants.py +++ b/bot/constants.py @@ -18,36 +18,10 @@ from pathlib import Path from typing import Dict, List import yaml -from yaml.constructor import ConstructorError log = logging.getLogger(__name__) -def _required_env_var_constructor(loader, node): - """ - Implements a custom YAML tag for loading required environment - variables. If the environment variable is set, this function - will simply return it. Otherwise, a `CRITICAL` log message is - given and the `KeyError` is re-raised. - - Example usage in the YAML configuration: - - bot: - token: !REQUIRED_ENV 'BOT_TOKEN' - """ - - value = loader.construct_scalar(node) - - try: - return os.environ[value] - except KeyError: - log.critical( - f"Environment variable `{value}` is required, but was not set. " - "Set it in your environment or override the option using it in your `config.yml`." - ) - raise - - def _env_var_constructor(loader, node): """ Implements a custom YAML tag for loading optional environment @@ -63,8 +37,12 @@ def _env_var_constructor(loader, node): default = None - try: - # Try to construct a list from this YAML node + # Check if the node is a plain string value + if node.id == 'scalar': + value = loader.construct_scalar(node) + key = str(value) + else: + # The node value is a list value = loader.construct_sequence(node) if len(value) >= 2: @@ -74,11 +52,6 @@ def _env_var_constructor(loader, node): else: # Otherwise, we just have a key key = value[0] - except ConstructorError: - # This YAML node is a plain value rather than a list, so we just have a key - value = loader.construct_scalar(node) - - key = str(value) return os.getenv(key, default) @@ -96,7 +69,9 @@ def _join_var_constructor(loader, node): yaml.SafeLoader.add_constructor("!ENV", _env_var_constructor) yaml.SafeLoader.add_constructor("!JOIN", _join_var_constructor) -yaml.SafeLoader.add_constructor("!REQUIRED_ENV", _required_env_var_constructor) + +# Pointing old tag to !ENV constructor to avoid breaking existing configs +yaml.SafeLoader.add_constructor("!REQUIRED_ENV", _env_var_constructor) with open("config-default.yml", encoding="UTF-8") as f: @@ -129,6 +104,34 @@ if Path("config.yml").exists(): _recursive_update(_CONFIG_YAML, user_config) +def check_required_keys(keys): + """ + Verifies that keys that are set to be required are present in the + loaded configuration. + """ + for key_path in keys: + lookup = _CONFIG_YAML + try: + for key in key_path.split('.'): + lookup = lookup[key] + if lookup is None: + raise KeyError(key) + except KeyError: + log.critical( + f"A configuration for `{key_path}` is required, but was not found. " + "Please set it in `config.yml` or setup an environment variable and try again." + ) + raise + + +try: + required_keys = _CONFIG_YAML['config']['required_keys'] +except KeyError: + pass +else: + check_required_keys(required_keys) + + class YAMLGetter(type): """ Implements a custom metaclass used for accessing diff --git a/config-default.yml b/config-default.yml index 2c814b11e..0019d1688 100644 --- a/config-default.yml +++ b/config-default.yml @@ -1,6 +1,6 @@ bot: - help_prefix: "bot." - token: !REQUIRED_ENV "BOT_TOKEN" + help_prefix: "bot." + token: !ENV "BOT_TOKEN" cooldowns: # Per channel, per tag. @@ -313,3 +313,6 @@ wolfram: user_limit_day: 10 guild_limit_day: 67 key: !ENV "WOLFRAM_API_KEY" + +config: + required_keys: ['bot.token'] |