tighten some generic types

Author: amorison

src/stagpy/commands.py

@@ -14,15 +14,16 @@
 from . import __version__, phyvars
 from ._helpers import baredoc, walk
 from .config import CONFIG_LOCAL, Config
-from .stagyydata import _sdat_from_conf
+from .stagyydata import StagyyData, _sdat_from_conf
 
 if typing.TYPE_CHECKING:
     from collections.abc import Iterable, Mapping
     from typing import Callable
 
     from loam.base import Section
 
-    from .datatypes import Varf, Varr, Vart
+    from .datatypes import Field, Rprof, Tseries, Varf, Varr, Vart
+    from .step import Step
 
 
 def info_cmd(conf: Config) -> None:
@@ -57,7 +58,10 @@ def info_cmd(conf: Config) -> None:
 
 def _layout(
     dict_vars: Mapping[str, Varf | Varr | Vart],
-    dict_vars_extra: Mapping[str, Callable],
+    dict_vars_extra: Mapping[
+        str,
+        Callable[[Step], Field | Rprof] | Callable[[StagyyData], Tseries],
+    ],
 ) -> Columns:
     """Print nicely [(var, description)] from phyvars."""
     desc = [(v, m.description) for v, m in dict_vars.items()]

src/stagpy/field.py

@@ -395,9 +395,9 @@ def cmd(conf: Config) -> None:
                         plot_vec(axis, step, var[1], conf=conf)
             if conf.field.timelabel:
                 time = step.timeinfo["t"]
-                time = _helpers.scilabel(time)
+                time_str = _helpers.scilabel(time)
                 axes[0, 0].text(
-                    0.02, 1.02, f"$t={time}$", transform=axes[0, 0].transAxes
+                    0.02, 1.02, f"$t={time_str}$", transform=axes[0, 0].transAxes
                 )
             oname = "_".join(chain.from_iterable(vfig))
             plt.tight_layout(w_pad=3)

src/stagpy/stagyydata.py

@@ -224,7 +224,7 @@ def isteps(self) -> NDArray[np.float64]:
         """
         return self._tseries.index.values
 
-    def at_step(self, istep: int) -> Series:
+    def at_step(self, istep: int) -> Series[np.float64]:
         """Time series output for a given step."""
         return self._tseries.loc[istep]  # type: ignore
 
@@ -650,7 +650,7 @@ class StagyyData:
             logic.
     """
 
-    path_hint: PathLike | str
+    path_hint: PathLike[str] | str
     read_parameters_dat: bool = True
 
     @property

src/stagpy/step.py

@@ -634,7 +634,7 @@ def geom(self) -> Geometry:
         return Geometry(self)
 
     @property
-    def timeinfo(self) -> Series:
+    def timeinfo(self) -> Series[np.float64]:
         """Time series data of the time step."""
         try:
             info = self.sdat.tseries.at_step(self.istep)