_load_pmg_settings look for .pmgrc.yaml in ~/.config first, then in ~

materialsproject · Oct 21, 2022 · 564671f · 564671f
1 parent 3e54dc9
commit 564671f
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 15 deletions.
diff --git a/pymatgen/cli/pmg_config.py b/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):

diff --git a/pymatgen/core/__init__.py b/pymatgen/core/__init__.py
@@ -28,29 +28,34 @@
 __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() -> dict[str, str]:
     settings = {}
     # Load .pmgrc.yaml file
+    yaml = YAML()
     try:
-        yaml = YAML()
         with open(SETTINGS_FILE) as yml_file:
             settings = yaml.load(yml_file)
     except FileNotFoundError:
-        pass
+        try:
+            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 k, v in os.environ.items():
-        if k.startswith("PMG_"):
-            settings[k] = v
-        elif k in ["VASP_PSP_DIR", "MAPI_KEY", "DEFAULT_FUNCTIONAL"]:
-            settings["PMG_" + k] = v
+    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