diff --git a/imas/_to_xarray.py b/imas/_to_xarray.py index 6caec501..13525c82 100644 --- a/imas/_to_xarray.py +++ b/imas/_to_xarray.py @@ -27,7 +27,7 @@ def to_xarray(ids: IDSToplevel, *paths: str) -> xarray.Dataset: # block checks if the paths are valid, and by using "metadata.path_string" we ensure # that / are used as separator. try: - paths = [ids.metadata[path].path_string for path in paths] + paths: list[str] = [ids.metadata[path].path_string for path in paths] except KeyError as exc: raise ValueError(str(exc)) from None diff --git a/imas/backends/imas_core/al_context.py b/imas/backends/imas_core/al_context.py index ede33bac..d3d2f620 100644 --- a/imas/backends/imas_core/al_context.py +++ b/imas/backends/imas_core/al_context.py @@ -115,7 +115,7 @@ def timerange_action( tmin: float, tmax: float, dtime: Optional[numpy.ndarray], - interpolation_method: int, + interpolation_method: Optional[int], ) -> "ALContext": """Begin a new timerange action for use in a ``with`` context.""" ctx = ll_interface.begin_timerange_action( @@ -163,7 +163,7 @@ def write_data(self, path: str, timebasepath: str, data: Any) -> None: """Call ual_write_data with this context.""" status = ll_interface.write_data(self.ctx, path, timebasepath, data) if status != 0: - raise LowlevelError(f"write data at {path!r}: {status=}") + raise LowlevelError(f"write data at {path!r}", status) def list_all_occurrences(self, ids_name: str) -> List[int]: """List all occurrences of this IDS.""" @@ -359,7 +359,7 @@ def timerange_action( tmin: float, tmax: float, dtime: Optional[numpy.ndarray], - interpolation_method: int, + interpolation_method: Optional[int], ) -> Iterator["LazyALContext"]: """Lazily start a lowlevel timerange action, see :meth:`ALContext.timerange_action`. diff --git a/imas/backends/imas_core/db_entry_al.py b/imas/backends/imas_core/db_entry_al.py index 8559b0c0..167d04b5 100644 --- a/imas/backends/imas_core/db_entry_al.py +++ b/imas/backends/imas_core/db_entry_al.py @@ -96,7 +96,6 @@ def from_pulse_run( options: Any, factory: IDSFactory, ) -> "ALDBEntryImpl": - # Set defaults user_name = user_name or getpass.getuser() data_version = data_version or factory.dd_version @@ -138,7 +137,7 @@ def _setup_backend(cls, backend: str, mode: int, factory: IDSFactory) -> None: if idsdef_path is None: # Extract XML from the DD zip and point UDA to it idsdef_path = extract_idsdef(factory.version) - os.environ["IDSDEF_PATH"] = idsdef_path + os.environ["IDSDEF_PATH"] = str(idsdef_path) elif backend in ["hdf5", "memory", "ascii", "flexbuffers"]: pass # nothing to set up @@ -173,7 +172,7 @@ def get( destination: IDSToplevel, lazy: bool, nbc_map: Optional[NBCPathMap], - ) -> None: + ) -> IDSToplevel: if self._db_ctx is None: raise RuntimeError("Database entry is not open.") if lazy and self.backend == "ascii": @@ -333,9 +332,12 @@ def delete_data(self, ids_name: str, occurrence: int) -> None: ll_path += f"/{occurrence}" ids = self._ids_factory.new(ids_name) with self._db_ctx.global_action(ll_path, WRITE_OP) as write_ctx: - delete_children(ids.metadata, write_ctx, "") + delete_children(ids.metadata, write_ctx) def list_all_occurrences(self, ids_name: str) -> List[int]: + if self._db_ctx is None: + raise RuntimeError("Database entry is not open.") + try: occurrence_list = self._db_ctx.list_all_occurrences(ids_name) except LLInterfaceError: diff --git a/imas/backends/imas_core/mdsplus_model.py b/imas/backends/imas_core/mdsplus_model.py index d5f667ee..df813e8e 100644 --- a/imas/backends/imas_core/mdsplus_model.py +++ b/imas/backends/imas_core/mdsplus_model.py @@ -366,7 +366,7 @@ def jTraverser_jar() -> Path: for component in os.environ.get("CLASSPATH", "").split(":"): if component.endswith(".jar"): if re.search(".*jTraverser.jar", component): - return component + return Path(component) else: # assume its a directory (strip any '*' suffix) search_dirs.append(component.rstrip("*")) diff --git a/imas/backends/imas_core/uda_support.py b/imas/backends/imas_core/uda_support.py index f051f549..299d6381 100644 --- a/imas/backends/imas_core/uda_support.py +++ b/imas/backends/imas_core/uda_support.py @@ -1,6 +1,6 @@ import logging from pathlib import Path -from typing import Union +from typing import Union, Optional from xml.etree import ElementTree as ET from imas import dd_zip @@ -10,7 +10,7 @@ logger = logging.getLogger(__name__) -def get_dd_version_from_idsdef_xml(path: Union[str, Path]) -> str: +def get_dd_version_from_idsdef_xml(path: Union[str, Path]) -> Optional[str]: """Parse the IDSDef.xml up to the point where the Data Dictionary version is set. Returns: diff --git a/imas/backends/netcdf/db_entry_nc.py b/imas/backends/netcdf/db_entry_nc.py index e6ee32cb..0776c47b 100644 --- a/imas/backends/netcdf/db_entry_nc.py +++ b/imas/backends/netcdf/db_entry_nc.py @@ -100,7 +100,7 @@ def get( destination: IDSToplevel, lazy: bool, nbc_map: Optional[NBCPathMap], - ) -> None: + ) -> IDSToplevel: # Feature compatibility checks if parameters is not None: if isinstance(parameters, GetSliceParameters): diff --git a/imas/backends/netcdf/nc_metadata.py b/imas/backends/netcdf/nc_metadata.py index 94929957..246376b2 100644 --- a/imas/backends/netcdf/nc_metadata.py +++ b/imas/backends/netcdf/nc_metadata.py @@ -1,7 +1,6 @@ # This file is part of IMAS-Python. # You should have received the IMAS-Python LICENSE file with this project. -"""NetCDF metadata for dimensions and tensorization of IDSs. -""" +"""NetCDF metadata for dimensions and tensorization of IDSs.""" from functools import lru_cache from typing import Dict, List, Optional, Set, Tuple @@ -89,7 +88,7 @@ def __init__(self, ids_metadata: IDSMetadata) -> None: # Add cache for public API self.get_dimensions = lru_cache(maxsize=None)(self.get_dimensions) - def get_coordinates(self, path: str, homogeneous_time: bool) -> Tuple[str]: + def get_coordinates(self, path: str, homogeneous_time: bool) -> Tuple[str, ...]: """Get the coordinates (adhering to CF conventions) for a netCDF variable. Args: @@ -109,7 +108,7 @@ def get_coordinates(self, path: str, homogeneous_time: bool) -> Tuple[str]: for coord in self.coordinates[path] ) - def get_dimensions(self, path: str, homogeneous_time: bool) -> Tuple[str]: + def get_dimensions(self, path: str, homogeneous_time: bool) -> Tuple[str, ...]: """Get the dimensions for a netCDF variable. Args: diff --git a/imas/ids_base.py b/imas/ids_base.py index 5c74bf5c..7d466fb8 100644 --- a/imas/ids_base.py +++ b/imas/ids_base.py @@ -1,7 +1,6 @@ # This file is part of IMAS-Python. # You should have received the IMAS-Python LICENSE file with this project. -"""Base class for all IDS nodes. -""" +"""Base class for all IDS nodes.""" import logging from typing import TYPE_CHECKING, Optional, Type @@ -12,6 +11,7 @@ if TYPE_CHECKING: from imas.ids_toplevel import IDSToplevel + from imas.ids_primitive import IDSInt0D logger = logging.getLogger(__name__) @@ -35,7 +35,7 @@ class IDSBase: """True iff this IDS lazy-loads its data""" @property - def _time_mode(self) -> int: + def _time_mode(self) -> "IDSInt0D": """Retrieve the time mode from `/ids_properties/homogeneous_time`""" return self._parent._time_mode diff --git a/imas/ids_convert.py b/imas/ids_convert.py index a1631b83..a1707c16 100644 --- a/imas/ids_convert.py +++ b/imas/ids_convert.py @@ -7,7 +7,7 @@ import logging from functools import lru_cache, partial from pathlib import Path -from typing import Callable, Dict, Iterator, List, Optional, Set, Tuple +from typing import Callable, Dict, Iterator, List, Optional, Set, Tuple, Any from xml.etree.ElementTree import Element, ElementTree import numpy @@ -70,7 +70,7 @@ def __init__(self) -> None: self.ctxpath: Dict[str, str] = {} """Map providing the lowlevel context path for renamed elements.""" - self.type_change: Dict[str, Optional[Callable[[IDSBase, IDSBase], None]]] = {} + self.type_change: Dict[str, Optional[Callable[[IDSBase, IDSBase], Any]]] = {} """Dictionary of paths that had a type change. Type changes are mapped to None in :py:attr:`path`, this ``dict`` allows to @@ -1001,9 +1001,7 @@ def _repeat_first_point(node: IDSBase) -> None: child.value = numpy.concatenate((child.value, [child.value[0]])) -def _remove_last_point_conditional( - source_node: IDSStructure, target_node: IDSStructure -) -> None: +def _remove_last_point_conditional(source_node: IDSBase, target_node: IDSBase) -> None: """Type change method for nbc_description=repeat_children_first_point_conditional*. This method handles converting from new (DDv4) to old (DDv3). diff --git a/imas/ids_metadata.py b/imas/ids_metadata.py index 4d2d5dbb..c1e29f43 100644 --- a/imas/ids_metadata.py +++ b/imas/ids_metadata.py @@ -1,7 +1,7 @@ # This file is part of IMAS-Python. # You should have received the IMAS-Python LICENSE file with this project. -"""Core of the IMAS-Python interpreted IDS metadata -""" +"""Core of the IMAS-Python interpreted IDS metadata""" + import re import types from enum import Enum @@ -77,7 +77,7 @@ def get_toplevel_metadata(structure_xml: Element) -> "IDSMetadata": IDSMetadata.__setattr__ = orig_setattr -_type_map: Dict[Tuple[IDSDataType, int], Type] = {} +_type_map: Dict[Tuple[Optional[IDSDataType], int], Type] = {} """Map of IDSDataType and ndim to IDSBase implementation class.""" @@ -205,11 +205,11 @@ def __init__( if self._parent is not None: self._is_dynamic = self.type.is_dynamic or self._parent._is_dynamic - self.coordinates: "tuple[IDSCoordinate]" + self.coordinates: "tuple[IDSCoordinate, ...]" """Tuple of coordinates of this node. ``coordinates[0]`` is the coordinate of the first dimension, etc.""" - self.coordinates_same_as: "tuple[IDSCoordinate]" + self.coordinates_same_as: "tuple[IDSCoordinate, ...]" """Indicates quantities which share the same coordinate in a given dimension, but the coordinate is not explicitly stored in the IDS.""" if self.ndim == 0: @@ -231,7 +231,7 @@ def __init__( self.coordinates_same_as = tuple(coors_same_as) # Parse alternative coordinates - self.alternative_coordinates: "tuple[IDSPath]" = () + self.alternative_coordinates: "tuple[IDSPath, ...]" = () """Quantities that can be used as coordinate instead of this node.""" if "alternative_coordinate1" in attrib: self.alternative_coordinates = tuple( diff --git a/imas/ids_toplevel.py b/imas/ids_toplevel.py index fcda5f0d..67e53c0f 100644 --- a/imas/ids_toplevel.py +++ b/imas/ids_toplevel.py @@ -25,6 +25,7 @@ ) from imas.ids_metadata import IDSMetadata, IDSType, get_toplevel_metadata from imas.ids_structure import IDSStructure +from imas.ids_primitive import IDSInt0D if TYPE_CHECKING: from imas.db_entry import DBEntry @@ -61,6 +62,7 @@ class IDSToplevel(IDSStructure): __doc__ = IDSDoc(__doc__) _path = "" # Path to ourselves without the IDS name and slashes + _parent: "IDSFactory" # In contrast to IDSBase, our parent is the IDSFactory def __init__(self, parent: "IDSFactory", structure_xml, lazy=False): """Save backend_version and backend_xml and build translation layer. @@ -89,7 +91,7 @@ def _dd_version(self) -> str: return self._version @property - def _time_mode(self) -> int: + def _time_mode(self) -> IDSInt0D: """Retrieve the time mode from `/ids_properties/homogeneous_time`""" return self.ids_properties.homogeneous_time diff --git a/tools/extract_test_data.py b/tools/extract_test_data.py deleted file mode 100644 index c17f8ec8..00000000 --- a/tools/extract_test_data.py +++ /dev/null @@ -1,57 +0,0 @@ -# This file is part of IMAS-Python. -# You should have received the IMAS-Python LICENSE file with this project. -import os - -import imas - -# Open input datafile -pulse, run, user, database = 134173, 106, "public", "ITER" -input = imas.DBEntry(imas.imasdef.MDSPLUS_BACKEND, database, pulse, run, user) -input.open() - -# Read Te profile and the associated normalised toroidal flux coordinate -get_these_idss = ["equilibrium", "core_profiles"] -idss = {} -# The reference has 871 timepoints -for time_index in [0, 433, 871]: - for ids_name in get_these_idss: - if ids_name not in idss: - idss[ids_name] = [] - idss[ids_name].append( - input.get_slice( - ids_name, - time_index, - imas.imasdef.PREVIOUS_INTERP, - occurrence=0, - ) - ) - -# Close the datafile -input.close() - -# Dump the data to ASCII -# Create output datafile -temp = imas.DBEntry(imas.imasdef.MEMORY_BACKEND, database, pulse, run, user) -temp.create() -for ids_name, ids_list in idss.items(): - for ids_slice in ids_list: - temp.put_slice(ids_slice) - -uber_idss = {} -for ids_name in idss: - uber_idss[ids_name] = temp.get(ids_name) -temp.close() - - -user = os.getenv("USER") -# Because we use the ASCII backend, this results in a .ids file in the cwd -output = imas.DBEntry(imas.imasdef.ASCII_BACKEND, database, pulse, run, user) -output.create() - -# Save the IDS -for ids_name, ids in uber_idss.items(): - print(f"Putting {ids_name}") - output.put(ids) - -# Close the output datafile -output.close()