Single source vars (#684)

* Load data vars from imas2xarray

* Reduce number of modules

* Bump version of imas2xarray

* Fix tests

pyproject.toml

@@ -38,7 +38,7 @@ dependencies = [
     "altair >= 5.0",
     "click",
     "gitpython",
-    "imas2xarray",
+    "imas2xarray >= 0.2.0",
     "importlib_resources; python_version<'3.10'",
     "jetto-tools >= 1.8.8",
     "jinja2 >= 3.0",

src/duqtools/config/_schema_root.py

@@ -3,17 +3,21 @@
 from pathlib import Path
 from typing import Optional, Union
 
+from imas2xarray import Variable, VariableConfigModel
 from pydantic import Field, PrivateAttr
 
 from duqtools.schema import BaseModel
 from duqtools.systems.ets import Ets6SystemModel
-from duqtools.systems.jetto import JettoSystemModel
+from duqtools.systems.jetto import IDS2JettoVariableModel, JettoSystemModel, JettoVariableModel
 from duqtools.systems.no_system import NoSystemModel
 
-from ._models import DuqtoolsVariableConfigModel
 from ._schema_create import CreateConfigModel
 
 
+class ExtraVariables(VariableConfigModel):
+    root: list[Union[JettoVariableModel, Variable, IDS2JettoVariableModel]]
+
+
 class ConfigModel(BaseModel):
     """The options for the CLI are defined by this model."""
     tag: str = Field(
@@ -26,7 +30,7 @@ class ConfigModel(BaseModel):
         description=
         'Configuration for the create subcommand. See model for more info.')
 
-    extra_variables: Optional[DuqtoolsVariableConfigModel] = Field(
+    extra_variables: Optional[ExtraVariables] = Field(
         None, description='Specify extra variables for this run.')
 
     system: Union[NoSystemModel, Ets6SystemModel, JettoSystemModel] = Field(

src/duqtools/config/_variables.py

@@ -1,22 +1,27 @@
 from __future__ import annotations
 
 import sys
+from typing import Union
 
-from imas2xarray import VariableConfigLoader
+from imas2xarray import VariableConfigLoader, VariableConfigModel
 
-from ._models import DuqtoolsVariableConfigModel
+from duqtools.systems.jetto import IDS2JettoVariableModel, JettoVariableModel
 
 if sys.version_info < (3, 10):
     from importlib_resources import files
 else:
     from importlib.resources import files
 
 
-class DuqtoolsVariableConfigLoader(VariableConfigLoader):
-    MODEL = DuqtoolsVariableConfigModel
-    VAR_DIR = 'duqtools'
-    VAR_ENV = 'DUQTOOLS_VARDEF'
-    MODULE = files('duqtools.data')
+class DuqtoolsVariableConfigModel(VariableConfigModel):
+    root: list[Union[JettoVariableModel, IDS2JettoVariableModel]]
 
 
-var_lookup = DuqtoolsVariableConfigLoader().load()
+imas2xarray_var_lookup = VariableConfigLoader().load()
+
+var_lookup = VariableConfigLoader(
+    model=DuqtoolsVariableConfigModel,
+    var_dir='duqtools',
+    var_env='DUQTOOLS_VARDEF',
+    module=files('duqtools.data'),
+).load(var_lookup=imas2xarray_var_lookup)

src/duqtools/list_variables.py

@@ -56,7 +56,7 @@ def list_variables(*, cfg: Config, **kwargs):
         list_group(group, extra_variables)
 
     grouped_other_vars = var_lookup.groupby_type()
-    grouped_other_vars.pop('IDS-variable')
+    grouped_other_vars.pop('IDS-variable', None)
 
     for var_type, group in grouped_other_vars.items():
         click.secho(f'\n{var_type}:', **ST_HEADER)