diff --git a/src/imas2xarray/_io.py b/src/imas2xarray/_io.py index 2d199da..b1acd64 100644 --- a/src/imas2xarray/_io.py +++ b/src/imas2xarray/_io.py @@ -238,7 +238,8 @@ def get_variables( """Get variables from data set. This function looks up the data location from the - `imas2xarray.var_lookup` table, and returns + `imas2xarray.var_lookup` table, and returns an xarray dataset. + Variable dimensions are automatically retrieved if available. Parameters ---------- @@ -263,8 +264,9 @@ def get_variables( """ var_models = var_lookup.lookup(variables) + # Attempt to automatically load associated dimensions dims = {dim for variable in var_models for dim in variable.dims} - var_models |= var_lookup.lookup(dims) + var_models |= var_lookup.lookup(dims, skip_missing=True) for var in var_models: if var.ids != ids: diff --git a/src/imas2xarray/_lookup.py b/src/imas2xarray/_lookup.py index 6264a76..4a0b00e 100644 --- a/src/imas2xarray/_lookup.py +++ b/src/imas2xarray/_lookup.py @@ -79,20 +79,37 @@ def groupby_ids(self) -> dict[Hashable, list[IDSVariableModel]]: grouped_ids_vars = groupby(ids_vars, keyfunc=lambda var: var.ids) return grouped_ids_vars - def lookup(self, variables: Collection[(str | IDSVariableModel)]) -> set[IDSVariableModel]: + def lookup( + self, variables: Collection[(str | IDSVariableModel)], skip_missing: bool = False + ) -> set[IDSVariableModel]: """Helper function to look up a bunch of variables. - If str, look up the variable from the `var_lookup`. Else, check if - the variable is an `IDSVariableModel`. + Parameters + ---------- + variables : Collection[(str | IDSVariableModel)] + List of variables to load. If str, look up the variable from the `var_lookup`. + Else, ensure the variable is an `IDSVariableModel`. + skip_missing : bool + Skip missing variables + + Returns + ------- + variables : set[IDSVariableModel] + Collection of variable models """ var_models = set() for var in variables: if isinstance(var, str): - if var.endswith(ERROR_SUFFIX): - var = self.error_upper(var) - else: - var = self[var] + try: + if var.endswith(ERROR_SUFFIX): + var = self.error_upper(var) + else: + var = self[var] + except KeyError: + if skip_missing: + continue + raise if not isinstance(var, IDSVariableModel): raise ValueError(f'Cannot lookup variable with type {type(var)}')