modernize annotations

Author: amorison

stagpy/_helpers.py

@@ -11,7 +11,7 @@
     from typing import Any, Optional
 
     from matplotlib.figure import Figure
-    from numpy import ndarray
+    from numpy.typing import NDArray
 
     from .config import Config
     from .stagyydata import StagyyData, _StepsView
@@ -102,7 +102,7 @@ def baredoc(obj: object) -> str:
     return doc.rstrip(" .").lstrip()
 
 
-def find_in_sorted_arr(value: Any, array: ndarray, after: bool = False) -> int:
+def find_in_sorted_arr(value: Any, array: NDArray, after: bool = False) -> int:
     """Return position of element in a sorted array.
 
     Returns:

stagpy/_step.py

@@ -20,19 +20,9 @@
 from .dimensions import Scales
 
 if typing.TYPE_CHECKING:
-    from typing import (
-        Any,
-        Callable,
-        Dict,
-        Iterator,
-        List,
-        Mapping,
-        NoReturn,
-        Optional,
-        Tuple,
-    )
-
-    from numpy import ndarray
+    from typing import Any, Callable, Iterator, Mapping, NoReturn, Optional
+
+    from numpy.typing import NDArray
     from pandas import DataFrame, Series
 
     from .datatypes import Varf
@@ -46,18 +36,18 @@ class _Geometry:
     output by StagYY.
     """
 
-    def __init__(self, header: Dict[str, Any], step: Step):
+    def __init__(self, header: dict[str, Any], step: Step):
         self._header = header
         self._step = step
-        self._shape: Dict[str, Any] = {
+        self._shape: dict[str, Any] = {
             "sph": False,
             "cyl": False,
             "axi": False,
             "ntot": list(header["nts"]) + [header["ntb"]],
         }
         self._init_shape()
 
-    def _scale_radius_mo(self, radius: ndarray) -> ndarray:
+    def _scale_radius_mo(self, radius: NDArray) -> NDArray:
         """Rescale radius for evolving MO runs."""
         if self._step.sdat.par.get("magma_oceans_in", "evolving_magma_oceans", False):
             return self._header["mo_thick_sol"] * (radius + self._header["mo_lambda"])
@@ -103,7 +93,7 @@ def nztot(self) -> int:
         return self.nrtot
 
     @cached_property
-    def r_walls(self) -> ndarray:
+    def r_walls(self) -> NDArray:
         """Position of FV walls along the z/r direction."""
         rgeom = self._header.get("rgeom")
         if rgeom is not None:
@@ -114,7 +104,7 @@ def r_walls(self) -> ndarray:
         return self._scale_radius_mo(walls)
 
     @cached_property
-    def r_centers(self) -> ndarray:
+    def r_centers(self) -> NDArray:
         """Position of FV centers along the z/r direction."""
         rgeom = self._header.get("rgeom")
         if rgeom is not None:
@@ -124,7 +114,7 @@ def r_centers(self) -> ndarray:
         return self._scale_radius_mo(walls)
 
     @cached_property
-    def t_walls(self) -> ndarray:
+    def t_walls(self) -> NDArray:
         """Position of FV walls along x/theta."""
         if self.threed or self.twod_xz:
             if self.yinyang:
@@ -143,12 +133,12 @@ def t_walls(self) -> ndarray:
         return np.array([center - d_t, center + d_t])
 
     @cached_property
-    def t_centers(self) -> ndarray:
+    def t_centers(self) -> NDArray:
         """Position of FV centers along x/theta."""
         return (self.t_walls[:-1] + self.t_walls[1:]) / 2
 
     @cached_property
-    def p_walls(self) -> ndarray:
+    def p_walls(self) -> NDArray:
         """Position of FV walls along y/phi."""
         if self.threed or self.twod_yz:
             if self.yinyang:
@@ -165,37 +155,37 @@ def p_walls(self) -> ndarray:
         return np.array([-d_p, d_p])
 
     @cached_property
-    def p_centers(self) -> ndarray:
+    def p_centers(self) -> NDArray:
         """Position of FV centers along y/phi."""
         return (self.p_walls[:-1] + self.p_walls[1:]) / 2
 
     @property
-    def z_walls(self) -> ndarray:
+    def z_walls(self) -> NDArray:
         """Same as r_walls."""
         return self.r_walls
 
     @property
-    def z_centers(self) -> ndarray:
+    def z_centers(self) -> NDArray:
         """Same as r_centers."""
         return self.r_centers
 
     @property
-    def x_walls(self) -> ndarray:
+    def x_walls(self) -> NDArray:
         """Same as t_walls."""
         return self.t_walls
 
     @property
-    def x_centers(self) -> ndarray:
+    def x_centers(self) -> NDArray:
         """Same as t_centers."""
         return self.t_centers
 
     @property
-    def y_walls(self) -> ndarray:
+    def y_walls(self) -> NDArray:
         """Same as p_walls."""
         return self.p_walls
 
     @property
-    def y_centers(self) -> ndarray:
+    def y_centers(self) -> NDArray:
         """Same as p_centers."""
         return self.p_centers
 
@@ -297,15 +287,15 @@ def __init__(
         step: Step,
         variables: Mapping[str, Varf],
         extravars: Mapping[str, Callable[[Step], Field]],
-        files: Mapping[str, List[str]],
-        filesh5: Mapping[str, List[str]],
+        files: Mapping[str, list[str]],
+        filesh5: Mapping[str, list[str]],
     ):
         self.step = step
         self._vars = variables
         self._extra = extravars
         self._files = files
         self._filesh5 = filesh5
-        self._data: Dict[str, Field] = {}
+        self._data: dict[str, Field] = {}
         super().__init__()
 
     def __getitem__(self, name: str) -> Field:
@@ -329,7 +319,7 @@ def __getitem__(self, name: str) -> Field:
         return self._data[name]
 
     @cached_property
-    def _present_fields(self) -> List[str]:
+    def _present_fields(self) -> list[str]:
         return [fld for fld in chain(self._vars, self._extra) if fld in self]
 
     def __iter__(self) -> Iterator[str]:
@@ -347,7 +337,7 @@ def __len__(self) -> int:
     def __eq__(self, other: object) -> bool:
         return self is other
 
-    def _get_raw_data(self, name: str) -> Tuple[List[str], Any]:
+    def _get_raw_data(self, name: str) -> tuple[list[str], Any]:
         """Find file holding data and return its content."""
         # try legacy first, then hdf5
         filestem = ""
@@ -384,7 +374,7 @@ def _get_raw_data(self, name: str) -> Tuple[List[str], Any]:
                     break
         return list_fvar, parsed_data
 
-    def _set(self, name: str, fld: ndarray) -> None:
+    def _set(self, name: str, fld: NDArray) -> None:
         sdat = self.step.sdat
         col_fld = sdat._collected_fields
         col_fld.append((self.step.istep, name))
@@ -399,7 +389,7 @@ def __delitem__(self, name: str) -> None:
             del self._data[name]
 
     @cached_property
-    def _header(self) -> Optional[Dict[str, Any]]:
+    def _header(self) -> Optional[dict[str, Any]]:
         if self.step.isnap is None:
             return None
         binfiles = self.step.sdat._binfiles_set(self.step.isnap)
@@ -439,9 +429,9 @@ class _Tracers:
 
     def __init__(self, step: Step):
         self.step = step
-        self._data: Dict[str, Optional[List[ndarray]]] = {}
+        self._data: dict[str, Optional[list[NDArray]]] = {}
 
-    def __getitem__(self, name: str) -> Optional[List[ndarray]]:
+    def __getitem__(self, name: str) -> Optional[list[NDArray]]:
         if name in self._data:
             return self._data[name]
         if self.step.isnap is None:
@@ -480,7 +470,7 @@ class _Rprofs:
 
     def __init__(self, step: Step):
         self.step = step
-        self._cached_extra: Dict[str, Rprof] = {}
+        self._cached_extra: dict[str, Rprof] = {}
 
     @cached_property
     def _data(self) -> Optional[DataFrame]:
@@ -527,12 +517,12 @@ def stepstr(self) -> str:
         return str(self.step.istep)
 
     @cached_property
-    def centers(self) -> ndarray:
+    def centers(self) -> NDArray:
         """Radial position of cell centers."""
         return self._rprofs["r"].to_numpy() + self.bounds[0]
 
     @cached_property
-    def walls(self) -> ndarray:
+    def walls(self) -> NDArray:
         """Radial position of cell walls."""
         rbot, rtop = self.bounds
         try:
@@ -547,7 +537,7 @@ def walls(self) -> ndarray:
         return walls
 
     @cached_property
-    def bounds(self) -> Tuple[float, float]:
+    def bounds(self) -> tuple[float, float]:
         """Radial or vertical position of boundaries.
 
         Radial/vertical positions of boundaries of the domain.

stagpy/args.py

@@ -26,7 +26,7 @@
 from .config import Config
 
 if typing.TYPE_CHECKING:
-    from typing import Any, Callable, List, Optional
+    from typing import Any, Callable, Optional
 
 
 def _sub(cmd: Any, *sections: str) -> Subcmd:
@@ -72,7 +72,7 @@ def _load_mplstyle(conf: Config) -> None:
 
 
 def parse_args(
-    conf: Config, arglist: Optional[List[str]] = None
+    conf: Config, arglist: Optional[list[str]] = None
 ) -> Callable[[Config], None]:
     """Parse cmd line arguments.
 

stagpy/commands.py

@@ -14,7 +14,7 @@
 from .config import CONFIG_LOCAL, Config
 
 if typing.TYPE_CHECKING:
-    from typing import Callable, Iterable, Mapping, Optional, Sequence, Tuple, Union
+    from typing import Callable, Iterable, Mapping, Optional, Sequence, Union
 
     from loam.base import Section
 
@@ -52,7 +52,7 @@ def info_cmd(conf: Config) -> None:
 
 
 def _pretty_print(
-    key_val: Sequence[Tuple[str, str]],
+    key_val: Sequence[tuple[str, str]],
     sep: str = ": ",
     min_col_width: int = 39,
     text_width: Optional[int] = None,

stagpy/field.py

@@ -17,13 +17,13 @@
 from .stagyydata import StagyyData
 
 if typing.TYPE_CHECKING:
-    from typing import Any, Dict, Iterable, Optional, Tuple, Union
+    from typing import Any, Iterable, Optional, Union
 
     from matplotlib.axes import Axes
     from matplotlib.collections import QuadMesh
     from matplotlib.colorbar import Colorbar
     from matplotlib.figure import Figure
-    from numpy import ndarray
+    from numpy.typing import NDArray
 
     from ._step import Step
     from .datatypes import Varf
@@ -36,7 +36,7 @@
 
 def _threed_extract(
     conf: Config, step: Step, var: str, walls: bool = False
-) -> Tuple[Tuple[ndarray, ndarray], ndarray]:
+) -> tuple[tuple[NDArray, NDArray], NDArray]:
     """Return suitable slices and coords for 3D fields."""
     is_vector = not valid_field_var(var)
     hwalls = is_vector or walls
@@ -89,7 +89,7 @@ def valid_field_var(var: str) -> bool:
 
 def get_meshes_fld(
     conf: Config, step: Step, var: str, walls: bool = False
-) -> Tuple[ndarray, ndarray, ndarray, Varf]:
+) -> tuple[NDArray, NDArray, NDArray, Varf]:
     """Return scalar field along with coordinates meshes.
 
     Only works properly in 2D geometry and 3D cartesian.
@@ -129,7 +129,7 @@ def get_meshes_fld(
 
 def get_meshes_vec(
     conf: Config, step: Step, var: str
-) -> Tuple[ndarray, ndarray, ndarray, ndarray]:
+) -> tuple[NDArray, NDArray, NDArray, NDArray]:
     """Return vector field components along with coordinates meshes.
 
     Only works properly in 2D geometry and 3D cartesian.
@@ -170,11 +170,11 @@ def get_meshes_vec(
 def plot_scalar(
     step: Step,
     var: str,
-    field: Optional[ndarray] = None,
+    field: Optional[NDArray] = None,
     axis: Optional[Axes] = None,
     conf: Optional[Config] = None,
     **extra: Any,
-) -> Tuple[Figure, Axes, QuadMesh, Optional[Colorbar]]:
+) -> tuple[Figure, Axes, QuadMesh, Optional[Colorbar]]:
     """Plot scalar field.
 
     Args:
@@ -264,7 +264,7 @@ def plot_iso(
     axis: Axes,
     step: Step,
     var: str,
-    field: Optional[ndarray] = None,
+    field: Optional[NDArray] = None,
     conf: Optional[Config] = None,
     **extra: Any,
 ) -> None:
@@ -289,7 +289,7 @@ def plot_iso(
 
     if conf.field.shift:
         fld = np.roll(fld, conf.field.shift, axis=0)
-    extra_opts: Dict[str, Any] = dict(linewidths=1)
+    extra_opts: dict[str, Any] = dict(linewidths=1)
     if "cmap" not in extra and conf.field.isocolors:
         extra_opts["colors"] = conf.field.isocolors
     elif "colors" not in extra:
@@ -337,9 +337,9 @@ def plot_vec(
 
 def _findminmax(
     view: _StepsView, sovs: Iterable[str]
-) -> Dict[str, Tuple[float, float]]:
+) -> dict[str, tuple[float, float]]:
     """Find min and max values of several fields."""
-    minmax: Dict[str, Tuple[float, float]] = {}
+    minmax: dict[str, tuple[float, float]] = {}
     for step in view.filter(snap=True):
         for var in sovs:
             if var in step.fields:
@@ -375,7 +375,7 @@ def cmd(conf: Config) -> None:
                 if var[0] not in step.fields:
                     print(f"{var[0]!r} field on snap {step.isnap} not found")
                     continue
-                opts: Dict[str, Any] = {}
+                opts: dict[str, Any] = {}
                 if var[0] in minmax:
                     opts = dict(vmin=minmax[var[0]][0], vmax=minmax[var[0]][1])
                 plot_scalar(step, var[0], axis=axis, conf=conf, **opts)