From 1bdca30be3b40d0328545fc8e5dc05ceb8eaaf6e Mon Sep 17 00:00:00 2001 From: Janosh Riebesell Date: Wed, 24 Jul 2024 17:17:22 -0400 Subject: [PATCH] Allow custom `.pmgrc.yaml` location via new `PMG_CONFIG_FILE` env var (#3949) --- src/pymatgen/core/__init__.py | 5 ++++- tests/core/test_settings.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/src/pymatgen/core/__init__.py b/src/pymatgen/core/__init__.py index c1f047de11c..848c251a8b4 100644 --- a/src/pymatgen/core/__init__.py +++ b/src/pymatgen/core/__init__.py @@ -37,9 +37,12 @@ def _load_pmg_settings() -> dict[str, Any]: settings: dict[str, Any] = {} + # PMG_CONFIG_FILE takes precedence over default settings location + settings_file = os.getenv("PMG_CONFIG_FILE") or SETTINGS_FILE + # Load .pmgrc.yaml file yaml = YAML() - for file_path in (SETTINGS_FILE, OLD_SETTINGS_FILE): + for file_path in (settings_file, OLD_SETTINGS_FILE): try: with open(file_path, encoding="utf-8") as yml_file: settings = yaml.load(yml_file) or {} diff --git a/tests/core/test_settings.py b/tests/core/test_settings.py index 39107becc90..6311909f6b1 100644 --- a/tests/core/test_settings.py +++ b/tests/core/test_settings.py @@ -48,3 +48,34 @@ def test_load_settings(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: # should return empty dict if file is invalid settings_file.write_text("---") assert _load_pmg_settings() == {} + + +def test_env_var_pmg_config_file(tmp_path: Path, monkeypatch: MonkeyPatch) -> None: + custom_config_file = tmp_path / "custom_config.yaml" + custom_config_file.write_text("PMG_CUSTOM_SETTING: custom_value") + + with monkeypatch.context() as ctx: + ctx.setenv("PMG_CONFIG_FILE", str(custom_config_file)) + settings = _load_pmg_settings() + assert "PMG_CUSTOM_SETTING" in settings + assert settings["PMG_CUSTOM_SETTING"] == "custom_value" + + # Test that PMG_CONFIG_FILE takes precedence over the default location + settings_file = tmp_path / ".pmgrc.yaml" + monkeypatch.setattr("pymatgen.core.SETTINGS_FILE", settings_file) + settings_file.write_text("PMG_DEFAULT_SETTING: default_value") + custom_config_file.write_text("PMG_CUSTOM_SETTING: custom_value") + + with monkeypatch.context() as ctx: + ctx.setenv("PMG_CONFIG_FILE", str(custom_config_file)) + settings = _load_pmg_settings() + assert "PMG_CUSTOM_SETTING" in settings + assert "PMG_DEFAULT_SETTING" not in settings + assert settings["PMG_CUSTOM_SETTING"] == "custom_value" + + # Test that env vars still take precedence over the values specified in PMG_CONFIG_FILE + with monkeypatch.context() as ctx: + ctx.setenv("PMG_CONFIG_FILE", str(custom_config_file)) + ctx.setenv("PMG_CUSTOM_SETTING", "env_value") + settings = _load_pmg_settings() + assert settings["PMG_CUSTOM_SETTING"] == "env_value"