flags unknown config keys and fixes a bug with env prefixes

src/quantem/core/config.py

@@ -40,6 +40,8 @@
 no_default = "__no_default__"
 config_lock = threading.Lock()
 
+ENV_PREFIX = "QUANTEM_"
+
 PATH = Path(os.getenv("QUANTEM_CONFIG", "~/.config/quantem")).expanduser().resolve()
 
 config: dict = {}
@@ -69,7 +71,7 @@ def __init__(
         **kwargs,
     ):
         self.config: dict = config
-        self._record = []
+        self._record: list[tuple[str, tuple[str, ...], Any]] = []
 
         if arg is not None:
             if not isinstance(arg, (Mapping)):
@@ -86,12 +88,23 @@ def __init__(
     def __enter__(self):
         return self.config
 
+    def __exit__(self, type, value, traceback):
+        for op, path, value in reversed(self._record):
+            d = self.config
+            for key in path[:-1]:
+                d = d[key]
+            if op == "replace":
+                d[path[-1]] = value
+            else:
+                d.pop(path[-1], None)
+
     def _assign(
         self,
         keys: Sequence[str],
         value: Any,
         d: dict,
         path: tuple[str, ...] = (),
+        record: bool = True,
     ) -> None:
         """Assign value into a nested configuration dictionary
 
@@ -105,17 +118,27 @@ def _assign(
             value
         path : tuple[str], optional
             The path history up to this point.
+        record : bool, optional
+            Whether this operation needs to be recorded to allow for rollback.
         """
         key = canonical_name(keys[0], d)
 
         path = path + (key,)
 
         if len(keys) == 1:
+            if record:
+                if key in d:
+                    self._record.append(("replace", path, d[key]))
+                else:
+                    self._record.append(("insert", path, None))
             d[key] = value
         else:
             if key not in d:
+                if record:
+                    self._record.append(("insert", path, None))
                 d[key] = {}
-            self._assign(keys[1:], value, d[key], path)
+                record = False
+            self._assign(keys[1:], value, d[key], path, record=record)
 
 
 def refresh(config: dict = config, defaults: list[Mapping] = defaults, **kwargs) -> None:
@@ -188,22 +211,24 @@ def update_defaults(new: dict, config: dict = config, defaults: list[Mapping] =
     2.  Updates the global config with the new configuration
         prioritizing older values over newer ones
     """
-    for key, value in new.items():
+    current_defaults = merge(*defaults)
+    # Registered before the keys are checked: they are what "known key" means.
+    defaults.append(new)
+
+    for key, value in list(new.items()):
         key, nval = check_key_val(key, value)
         new[key] = nval
 
-    current_defaults = merge(*defaults)
-    defaults.append(new)
     update(config, new, priority="new-defaults", defaults=current_defaults)
 
 
 def _initialize() -> None:
     fn = os.path.join(os.path.dirname(__file__), "quantem.yaml")
 
     with open(fn) as f:
-        _defaults = yaml.safe_load(f)
+        shipped = yaml.safe_load(f)
 
-    update_defaults(_defaults)
+    update_defaults(shipped)
 
 
 def canonical_name(k: str, config: dict) -> str:
@@ -234,6 +259,7 @@ def update(
     new: Mapping,
     priority: Literal["old", "new", "new-defaults"] = "new",
     defaults: Mapping | None = None,
+    check: bool = True,
 ) -> dict:
     """Update a nested dictionary with values from another
 
@@ -249,6 +275,9 @@ def update(
         If 'new-defaults', a mapping should be given of the current defaults.
         Only if a value in ``old`` matches the current default, it will be
         updated with ``new``.
+    check: bool
+        Whether to run the keys through :func:`check_key_val`. False on the
+        recursive call, because the unknown-key warning is about top-level keys.
 
     Examples
     --------
@@ -270,7 +299,8 @@ def update(
 
     """
     for k, v in new.items():
-        k, v = check_key_val(k, v)
+        if check:
+            k, v = check_key_val(k, v)
         k = canonical_name(k, old)
 
         if isinstance(v, Mapping):
@@ -281,6 +311,7 @@ def update(
                 v,
                 priority=priority,
                 defaults=defaults.get(k) if defaults else None,
+                check=False,
             )
         else:
             if (
@@ -308,7 +339,8 @@ def collect(path: Path | str = PATH, env: Mapping[str, str] | None = None) -> di
         A list of paths to search for yaml config files
 
     env : Mapping[str, str]
-        The system environment variables
+        The system environment variables. Values found here take
+        precedence over those read from the yaml files.
 
     Returns
     -------
@@ -318,8 +350,7 @@ def collect(path: Path | str = PATH, env: Mapping[str, str] | None = None) -> di
     if env is None:
         env = os.environ
 
-    # configs = [*collect_yaml(paths=paths), collect_env(env=env)] # skipping env
-    configs = list([*collect_yaml(path=Path(path))])
+    configs = [*collect_yaml(path=Path(path)), collect_env(env=env)]
     return merge(*configs)
 
 
@@ -358,7 +389,7 @@ def collect_env(env: Mapping[str, str] | None = None) -> dict:
     turns these into config variables of the form ``{"foo": {"bar-baz": 123}}``
     It transforms the key and value in the following way:
 
-    -  Lower-cases the key text
+    -  Strips the ``QUANTEM_`` prefix and lower-cases the rest
     -  Treats ``__`` (double-underscore) as nested access
     -  Calls ``ast.literal_eval`` on the value
     """
@@ -369,8 +400,10 @@ def collect_env(env: Mapping[str, str] | None = None) -> dict:
     d = {}
 
     for name, value in env.items():
-        if name.startswith("QUANTEM_"):
-            varname = name[5:].lower().replace("__", ".")
+        # QUANTEM_CONFIG says where the config files are; it is not one of
+        # the keys they hold.
+        if name.startswith(ENV_PREFIX) and name != "QUANTEM_CONFIG":
+            varname = name[len(ENV_PREFIX) :].lower().replace("__", ".")
             d[varname] = interpret_value(value)
 
     result: dict = {}
@@ -403,7 +436,7 @@ def merge(*dicts: Mapping) -> dict:
     """
     result: dict = {}
     for d in dicts:
-        update(result, d)
+        update(result, d, check=False)
     return result
 
 
@@ -429,12 +462,17 @@ def _load_config_file(path: str) -> dict | None:
 
 
 def check_key_val(key: str, val: Any, deprecations: dict = deprecations) -> tuple[str, Any]:
-    """Check if the provided value has been renamed or removed
+    """Check whether a key has been renamed, removed, or is not one we ship
+
+    A key that is none of the registered defaults warns and is still set: config
+    is not a schema, and refusing an unknown key would break anything that stores
+    its own.
 
     Parameters
     ----------
     key : str
-        The configuration key to check
+        The configuration key to check. May be dotted, in which case only the
+        part before the first '.' is checked.
     deprecations : Dict[str, str]
         The mapping of aliases
 
@@ -470,6 +508,13 @@ def check_key_val(key: str, val: Any, deprecations: dict = deprecations) -> tupl
         else:
             raise ValueError(f'Configuration value "{key}" has been removed')
 
+    top = key.split(".")[0]
+    # The top-level keys of the registered defaults, not merge(*defaults):
+    # merge() goes through update(), which calls back into here.
+    known = {k for d in defaults for k in d}
+    if top not in known and top not in deprecations:
+        warnings.warn(f'Unknown configuration key "{key}"')
+
     new_val = val
     if key in aliases:
         val_aliases = aliases[key]
@@ -594,5 +639,5 @@ def device() -> str:
     return get("device")
 
 
-refresh()
 _initialize()
+refresh()

src/quantem/core/quantem.yaml

@@ -1,13 +1,10 @@
-# Defualt configuration file for quantem
-# TODO: control for threads, torch
+# Default configuration file for quantem
 
 # The device can be either 'cpu', "gpu", or "cuda:<x>". "gpu" defaults to "cuda:0"
 device: cpu
 
 # The float precision to use. Options are 'float32' or 'float64'
-# ?? will this control both real and complex?
 precision: float32
-# global verbosity, ?? unecessary
 
 dtype_real: float32
 dtype_complex: complex64
@@ -21,7 +18,7 @@ cupy:
 mkl:
   # The number of threads to use for mkl
   threads: 2
-warnings: ### not sure how handling warnings yet
+warnings:
 
   # Show a warning when the grid is overspecified
 
@@ -34,20 +31,13 @@ viz:
   # The default units to use in real space
   real_space_units: "A"
   # The default units to use in reciprocal space
-
-  reciprocal_space_units: "Angstrom"
-  # The default colormap to use for plotting
-  cmap: "viridis"
-  # The default colormap to use for plotting the phase
-  phase_cmap: "hsluv"
-
-  default_colors: ""
-
   reciprocal_space_units: "A^-1"
   # The default colormap to use for showing images
   cmap: "gray"
   # The default colormap to use for showing the phase
   phase_cmap: "magma"
+
+  default_colors: ""
   # plotting default colors: 
   colors: 
     # categorical color lists for plotting

tests/test_config.py

@@ -1,3 +1,5 @@
+import warnings
+
 import pytest
 import torch
 
@@ -47,3 +49,102 @@ def test_config_update_defaults():
     config.refresh()
     assert config.get("dtype_real") == "int32"
     config.update_defaults({"dtype_real": start_dtype})
+
+
+def test_collect_env_flat_and_nested():
+    collected = config.collect_env(
+        env={
+            "QUANTEM_VERBOSE": "3",
+            "QUANTEM_VIZ__CMAP": "magma",
+            "PATH": "/usr/bin",
+        }
+    )
+    assert collected == {"verbose": 3, "viz": {"cmap": "magma"}}
+
+
+def test_collect_env_interprets_values():
+    collected = config.collect_env(
+        env={
+            "QUANTEM_VERBOSE": "3",
+            "QUANTEM_HAS_CUPY": "False",
+            "QUANTEM_VIZ__DEFAULT_COLORS": "['#000000', '#ffffff']",
+        }
+    )
+    assert collected["verbose"] == 3
+    assert collected["has_cupy"] is False
+    assert collected["viz"]["default_colors"] == ["#000000", "#ffffff"]
+
+
+def test_collect_env_ignores_quantem_config():
+    """QUANTEM_CONFIG names the config directory, it is not a config key."""
+    assert config.collect_env(env={"QUANTEM_CONFIG": "/some/dir"}) == {}
+
+
+def test_collect_env_overrides_yaml(tmp_path):
+    (tmp_path / "quantem.yaml").write_text("verbose: 5\nviz:\n  cmap: viridis\n")
+
+    from_yaml = config.collect(path=tmp_path, env={})
+    assert from_yaml == {"verbose": 5, "viz": {"cmap": "viridis"}}
+
+    merged = config.collect(path=tmp_path, env={"QUANTEM_VERBOSE": "9"})
+    assert merged == {"verbose": 9, "viz": {"cmap": "viridis"}}
+
+
+def test_check_key_val_known_key_silent():
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        config.check_key_val("verbose", 1)
+    assert caught == []
+
+
+def test_check_key_val_nested_key_silent():
+    """Only the top-level key is checked, a key under it is not."""
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        config.check_key_val("viz.not_a_real_key", 1)
+    assert caught == []
+
+
+def test_check_key_val_unknown_key_warns_and_still_sets():
+    with pytest.warns(UserWarning, match='Unknown configuration key "not_a_real_key"'):
+        with config.set({"not_a_real_key": 7}):
+            assert config.get("not_a_real_key") == 7
+    assert config.get("not_a_real_key", None) is None
+
+
+def test_check_key_val_deprecated_key_warns_as_deprecated():
+    config.deprecations["old_key"] = "verbose"
+    try:
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            config.check_key_val("old_key", 1)
+        assert len(caught) == 1
+        assert "has been deprecated" in str(caught[0].message)
+    finally:
+        del config.deprecations["old_key"]
+
+
+def test_set_context_manager_rolls_back():
+    d = {"verbose": 1}
+    with config.set({"verbose": 2, "precision": "float64"}, config=d):
+        assert d == {"verbose": 2, "precision": "float64"}
+    assert d == {"verbose": 1}
+
+
+def test_set_context_manager_rolls_back_nested():
+    d = {"verbose": 1}
+    with config.set({"viz.cmap": "magma"}, config=d):
+        assert d == {"verbose": 1, "viz": {"cmap": "magma"}}
+    assert d == {"verbose": 1}
+
+    d = {"viz": {"cmap": "gray"}}
+    with config.set({"viz.cmap": "magma"}, config=d):
+        assert d == {"viz": {"cmap": "magma"}}
+    assert d == {"viz": {"cmap": "gray"}}
+
+
+def test_set_context_manager_rolls_back_global_config():
+    start = config.get("verbose")
+    with config.set({"verbose": 42}):
+        assert config.get("verbose") == 42
+    assert config.get("verbose") == start