satiate mypy

pycontrails/models/ps_model/ps_aircraft_params.py

@@ -13,6 +13,7 @@
 
 from pycontrails.core.aircraft_performance import AircraftPerformanceParams
 from pycontrails.physics import constants as c
+from pycontrails.utils.types import ArrayOrFloat
 
 #: Path to the Poll-Schumann aircraft parameters CSV file.
 PS_FILE_PATH = pathlib.Path(__file__).parent / "static" / "ps-aircraft-params-20240524.csv"
@@ -260,18 +261,18 @@ def load_aircraft_engine_params(
     return dict(_row_to_aircraft_engine_params(tup) for tup in df.itertuples(index=False))
 
 
-def turbine_entry_temperature_at_max_take_off(first_flight: float) -> float:
+def turbine_entry_temperature_at_max_take_off(first_flight: ArrayOrFloat) -> ArrayOrFloat:
     """
     Calculate turbine entry temperature at maximum take-off rating.
 
     Parameters
     ----------
-    first_flight: float
+    first_flight: ArrayOrFloat
         Year of first flight
 
     Returns
     -------
-    float
+    ArrayOrFloat
         Turbine entry temperature at maximum take-off rating, ``tet_mto``, [:math:`K`]
 
     Notes
@@ -285,7 +286,7 @@ def turbine_entry_temperature_at_max_take_off(first_flight: float) -> float:
     - :cite:`cumpstyJetPropulsion2015`
     """
     out = 2000.0 * (1.0 - np.exp(62.8 - 0.0325 * first_flight))
-    if isinstance(out, np.ndarray):
+    if isinstance(first_flight, np.ndarray):
         return out
     return out.item()
 

pycontrails/models/ps_model/ps_grid.py

@@ -309,7 +309,6 @@ def _parse_variables(
 
     Returns a tuple of ``(dims, coords, air_pressure, air_temperature)``.
     """
-
     if isinstance(air_temperature, xr.DataArray):
         if level is not None:
             msg = "If 'air_temperature' is a DataArray, 'level' must be None"
@@ -322,21 +321,22 @@ def _parse_variables(
             raise KeyError(msg) from exc
 
         air_temperature, pressure_da = xr.broadcast(air_temperature, pressure_da)
-        dims = air_temperature.dims
-        coords = air_temperature.coords
-        return dims, coords, np.asarray(pressure_da), np.asarray(air_temperature)
+        return (  # type: ignore[return-value]
+            air_temperature.dims,
+            air_temperature.coords,
+            np.asarray(pressure_da),
+            np.asarray(air_temperature),
+        )
 
     if level is None:
         msg = "The 'level' argument must be provided"
         raise ValueError(msg)
 
-    dims = ("level",)
-    coords = {"level": level}
     air_pressure = level * 100.0
     if air_temperature is None:
         altitude_m = units.pl_to_m(level)
         air_temperature = units.m_to_T_isa(altitude_m)
-    return dims, coords, air_pressure, air_temperature
+    return ("level",), {"level": level}, air_pressure, air_temperature
 
 
 def ps_nominal_grid(