diff --git a/src/quantem/core/config.py b/src/quantem/core/config.py index c2eda10ba..3eae4faaf 100644 --- a/src/quantem/core/config.py +++ b/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,12 +211,14 @@ 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) @@ -201,9 +226,9 @@ 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() diff --git a/src/quantem/core/quantem.yaml b/src/quantem/core/quantem.yaml index 80f34587a..033d35737 100644 --- a/src/quantem/core/quantem.yaml +++ b/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:". "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 diff --git a/tests/test_config.py b/tests/test_config.py index 975a0613d..7dc1037c0 100644 --- a/tests/test_config.py +++ b/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