Merge pull request #2698 from materialsproject/env-vars-override-pmgrc

Let env vars override `.pmgrc.yaml` in `_load_pmg_settings()` and make `~/.config` new default location for `.pmgrc.yaml`

pymatgen/cli/pmg_config.py

@@ -19,7 +19,7 @@
 
 from monty.serialization import dumpfn, loadfn
 
-from pymatgen.core import SETTINGS_FILE
+from pymatgen.core import OLD_SETTINGS_FILE, SETTINGS_FILE
 
 
 def setup_potcars(potcar_dirs: list[str]):
@@ -178,18 +178,27 @@ def install_software(install: Literal["enumlib", "bader"]):
 
 def add_config_var(tokens: list[str], backup_suffix: str) -> None:
     """Add/update keys in .pmgrc.yaml config file."""
-    d = {}
     if os.path.exists(SETTINGS_FILE):
+        # read and write new config file if exists
+        fpath = SETTINGS_FILE
+    elif os.path.exists(OLD_SETTINGS_FILE):
+        # else use old config file if exists
+        fpath = OLD_SETTINGS_FILE
+    else:
+        # if neither exists, create new config file
+        fpath = SETTINGS_FILE
+    d = {}
+    if os.path.exists(fpath):
         if backup_suffix:
-            shutil.copy(SETTINGS_FILE, SETTINGS_FILE + backup_suffix)
-            print(f"Existing {SETTINGS_FILE} backed up to {SETTINGS_FILE}{backup_suffix}")
-        d = loadfn(SETTINGS_FILE)
+            shutil.copy(fpath, fpath + backup_suffix)
+            print(f"Existing {fpath} backed up to {fpath}{backup_suffix}")
+        d = loadfn(fpath)
     if len(tokens) % 2 != 0:
         raise ValueError(f"Uneven number {len(tokens)} of tokens passed to pmg config. Needs a value for every key.")
     for key, val in zip(tokens[0::2], tokens[1::2]):
         d[key] = val
-    dumpfn(d, SETTINGS_FILE)
-    print(f"New {SETTINGS_FILE} written!")
+    dumpfn(d, fpath)
+    print(f"New {fpath} written!")
 
 
 def configure_pmg(args: Namespace):

pymatgen/core/__init__.py

@@ -6,6 +6,8 @@
 operations on them.
 """
 
+from __future__ import annotations
+
 import os
 import warnings
 
@@ -26,31 +28,36 @@
 __version__ = "2022.9.21"
 
 
-SETTINGS_FILE = os.path.join(os.path.expanduser("~"), ".pmgrc.yaml")
-
+SETTINGS_FILE = os.path.join(os.path.expanduser("~"), ".config", ".pmgrc.yaml")
+OLD_SETTINGS_FILE = os.path.join(os.path.expanduser("~"), ".pmgrc.yaml")
 
-def _load_pmg_settings():
-    # Load environment variables by default as backup
-    d = {}
-    for k, v in os.environ.items():
-        if k.startswith("PMG_"):
-            d[k] = v
-        elif k in ["VASP_PSP_DIR", "MAPI_KEY", "DEFAULT_FUNCTIONAL"]:
-            d["PMG_" + k] = v
 
-    # Override anything in env vars with that in yml file
-    if os.path.exists(SETTINGS_FILE):
+def _load_pmg_settings() -> dict[str, str]:
+    settings = {}
+    # Load .pmgrc.yaml file
+    yaml = YAML()
+    try:
+        with open(SETTINGS_FILE) as yml_file:
+            settings = yaml.load(yml_file)
+    except FileNotFoundError:
         try:
-            yaml = YAML()
-            with open(SETTINGS_FILE) as f:
-                d_yml = yaml.load(f)
-            d.update(d_yml)
-        except Exception as ex:
-            # If there are any errors, default to using environment variables
-            # if present.
-            warnings.warn(f"Error loading .pmgrc.yaml: {ex}. You may need to reconfigure your yaml file.")
-
-    return d
+            with open(OLD_SETTINGS_FILE) as yml_file:
+                settings = yaml.load(yml_file)
+        except FileNotFoundError:
+            pass
+    except Exception as ex:
+        # If there are any errors, default to using environment variables
+        # if present.
+        warnings.warn(f"Error loading .pmgrc.yaml: {ex}. You may need to reconfigure your yaml file.")
+
+    # Override .pmgrc.yaml with env vars if present
+    for key, val in os.environ.items():
+        if key.startswith("PMG_"):
+            settings[key] = val
+        elif key in ("VASP_PSP_DIR", "MAPI_KEY", "DEFAULT_FUNCTIONAL"):
+            settings[f"PMG_{key}"] = val
+
+    return settings
 
 
 SETTINGS = _load_pmg_settings()

pymatgen/io/vasp/inputs.py

@@ -470,7 +470,6 @@ def get_string(self, direct: bool = True, vasp4_compatible: bool = False, signif
         Returns:
             String representation of POSCAR.
         """
-
         # This corrects for VASP really annoying bug of crashing on lattices
         # which have triple product < 0. We will just invert the lattice
         # vectors.
@@ -1244,12 +1243,9 @@ def automatic_gamma_density(structure: Structure, kppa: float):
             reciprocal lattice vector proportional to its length.
 
         Args:
-            structure:
-                Input structure
-            kppa:
-                Grid density
+            structure: Input structure
+            kppa: Grid density
         """
-
         latt = structure.lattice
         lengths = latt.abc
         ngrid = kppa / structure.num_sites
@@ -1889,8 +1885,8 @@ def from_symbol_and_functional(symbol: str, functional: str = None):
         :param functional: E.g., PBE
         :return: PotcarSingle
         """
-        if functional is None:
-            functional = SETTINGS.get("PMG_DEFAULT_FUNCTIONAL", "PBE")
+        functional = functional or SETTINGS.get("PMG_DEFAULT_FUNCTIONAL", "PBE")
+        assert isinstance(functional, str)  # type narrowing
         funcdir = PotcarSingle.functional_dir[functional]
         d = SETTINGS.get("PMG_VASP_PSP_DIR")
         if d is None:
@@ -2158,7 +2154,6 @@ def __getattr__(self, a):
         For float type properties, they are converted to the correct float. By
         default, all energies in eV and all length scales are in Angstroms.
         """
-
         try:
             return self.keywords[a.upper()]
         except Exception:
@@ -2420,7 +2415,9 @@ def run_vasp(
         :param err_file: File to write err.
         """
         self.write_input(output_dir=run_dir)
-        vasp_cmd = vasp_cmd or SETTINGS.get("PMG_VASP_EXE")
+        vasp_cmd = vasp_cmd or SETTINGS.get("PMG_VASP_EXE")  # type: ignore[assignment]
+        if not vasp_cmd:
+            raise ValueError("No VASP executable specified!")
         vasp_cmd = [os.path.expanduser(os.path.expandvars(t)) for t in vasp_cmd]
         if not vasp_cmd:
             raise RuntimeError("You need to supply vasp_cmd or set the PMG_VASP_EXE in .pmgrc.yaml to run VASP.")