From aa7087c5c1a46adf02cac7c8b6e66f8b9e309ed9 Mon Sep 17 00:00:00 2001 From: Stef Smeets Date: Tue, 5 Dec 2023 16:29:04 +0100 Subject: [PATCH 1/2] Specify search locations as argument --- src/imas2xarray/_lookup.py | 53 +++++++++++++++++++++++++------------- 1 file changed, 35 insertions(+), 18 deletions(-) diff --git a/src/imas2xarray/_lookup.py b/src/imas2xarray/_lookup.py index 6e0d6a4..b08bcd5 100644 --- a/src/imas2xarray/_lookup.py +++ b/src/imas2xarray/_lookup.py @@ -6,7 +6,7 @@ import sys from collections import UserDict from pathlib import Path, PosixPath -from typing import Hashable, Sequence +from typing import Any, Hashable, Sequence from pydantic_yaml import parse_yaml_raw_as @@ -22,7 +22,6 @@ USER_CONFIG_HOME = Path.home() / '.config' LOCAL_DIR = Path('.').absolute() -VAR_FILENAME = 'variables.yaml' VAR_FILENAME_GLOB = 'variables*.yaml' ERROR_SUFFIX = '_error_upper' @@ -99,22 +98,40 @@ def lookup(self, variables: Sequence[(str | IDSVariableModel)]) -> list[IDSVaria class VariableConfigLoader: - MODEL = VariableConfigModel - VAR_DIR = 'imas2xarray' - VAR_ENV = 'IMAS2XARRAY_VARDEF' - MODULE = files('imas2xarray.data') + def __init__( + self, + *, + model: type = VariableConfigModel, + var_dir: str = 'imas2xarray', + var_env: str = 'IMAS2XARRAY_VARDEF', + module: Path | Any = files('imas2xarray.data'), + ): + self.model = model + self.var_dir = var_dir + self.var_env = var_env + self.module = module - def __init__(self): self.paths = self.get_config_path() - def load(self) -> VarLookup: - """Load the variables config.""" - var_lookup = VarLookup() + def load(self, var_lookup: None | VarLookup = None) -> VarLookup: + """Load the variables config. + + Parameters + ---------- + var_lookup : None | VarLookup + Populate variable lookup table to initialize list of variables. Use + this to load variables from different locations + + Returns: + VarLookup: Description + """ + if not var_lookup: + var_lookup = VarLookup() for path in self.paths: logger.debug(f'Loading variables from: {path}') with open(path) as f: - var_config = parse_yaml_raw_as(self.MODEL, f) + var_config = parse_yaml_raw_as(self.model, f) var_lookup.update(var_config.to_variable_dict()) return var_lookup @@ -139,13 +156,13 @@ def get_config_path(self) -> tuple[Path, ...]: return self._get_paths_fallback() def _get_paths_from_environment_variable(self) -> tuple[Path, ...] | None: - env = os.environ.get(self.VAR_ENV) + env = os.environ.get(self.var_env) if env: path = Path(env) drc = path.parent if not drc.exists(): - raise OSError(f'{path} defined by ${self.VAR_ENV} does not exist!') + raise OSError(f'{path} defined by ${self.var_env} does not exist!') return tuple(drc.glob(path.name)) @@ -157,19 +174,19 @@ def _get_paths_local_directory(self) -> tuple[Path, ...] | None: def _get_paths_from_config_home(self) -> tuple[Path, ...] | None: config_home = os.environ.get('XDG_CONFIG_HOME', USER_CONFIG_HOME) - drc = Path(config_home) / self.VAR_DIR + drc = Path(config_home) / self.var_dir if drc.exists(): return tuple(drc.glob(VAR_FILENAME_GLOB)) return None def _get_paths_fallback(self) -> tuple[Path, ...]: - assert self.MODULE.is_dir() + assert self.module.is_dir() - if isinstance(self.MODULE, PosixPath): - drc = self.MODULE + if isinstance(self.module, PosixPath): + drc = self.module else: - drc = self.MODULE._paths[0] # type: ignore + drc = self.module._paths[0] # type: ignore return tuple(drc.glob(VAR_FILENAME_GLOB)) From d5a8591956402a83f90af3e4753797a94517b9c8 Mon Sep 17 00:00:00 2001 From: Stef Smeets Date: Tue, 5 Dec 2023 16:29:13 +0100 Subject: [PATCH 2/2] Fix typo --- docs/plotting_example.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/plotting_example.md b/docs/plotting_example.md index d300257..53991db 100644 --- a/docs/plotting_example.md +++ b/docs/plotting_example.md @@ -10,7 +10,7 @@ Below is an example of how to use **imas2xarray** to plot data with [matplotlib] The code below shows how to make a plot with [matplotlib](https://matplotlib.org/) for multiple datasets. -For a more advanced example of how to concatenate data, check out the [example notebooks../notebooks/xarray). +For a more advanced example of how to concatenate data, check out the [example notebooks](../notebooks/xarray). ```python {!../scripts/plot_with_matplotlib_multi.py!}