Fix typechecking errors (#2554)

* typing

* ignore

* typing

* typing

Author: VeckoTheGecko

pyproject.toml

@@ -161,4 +161,5 @@ known-first-party = ["parcels"]
 include = ["./src/"]
 exclude = [
   "./src/parcels/interpolators/", # ignore for now
+  "./src/parcels/kernels/_advection.py",
 ]

src/parcels/_core/basegrid.py

@@ -113,7 +113,7 @@ def ravel_index(self, axis_indices: dict[str, np.ndarray]) -> np.ndarray:
         indices = np.array([axis_indices[axis] for axis in self.axes], dtype=int)
         return _ravel(dims, indices)
 
-    def unravel_index(self, ei: int) -> dict[str, int]:
+    def unravel_index(self, ei: int) -> dict[str, np.ndarray]:
         """
         Convert a single encoded index (ei) back to a dictionary of axis indices.
 

src/parcels/_core/index_search.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from datetime import datetime
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -63,7 +62,7 @@ def _search_1d_array(
     return np.atleast_1d(index), np.atleast_1d(bcoord)
 
 
-def _search_time_index(field: Field, time: datetime):
+def _search_time_index(field: Field, time: np.ndarray):
     """Find and return the index and relative coordinate in the time array associated with a given time.
 
     Parameters

src/parcels/_core/utils/time.py

@@ -58,7 +58,7 @@ def time_length_as_flt(self):
     def __contains__(self, item: T) -> bool:
         return self.left <= item <= self.right
 
-    def is_all_time_in_interval(self, time: float):
+    def is_all_time_in_interval(self, time: float | np.ndarray):
         item = np.atleast_1d(time)
         return (0 <= item).all() and (item <= self.time_length_as_flt).all()
 

src/parcels/_core/xgrid.py

@@ -394,7 +394,7 @@ def get_xgcm_position_from_dim_name(axes: ptyping.XgcmAxes, dim: str) -> ptyping
         var_to_position = {var: position for position, var in axis.coords.items()}
 
         if dim in var_to_position:
-            return var_to_position[dim]
+            return var_to_position[dim]  # type: ignore[invalid-return-type] # due to mistyping in xgcm
     return None
 
 

src/parcels/_datasets/utils.py

@@ -206,7 +206,7 @@ def from_xarray_dataset_dict(d) -> xr.Dataset:
 def _fill_with_dummy_data(d: dict[str, dict]):
     assert isinstance(d, dict)
     if "dtype" in d:
-        d["data"] = np.zeros(d["shape"], dtype=d["dtype"])
+        d["data"] = np.zeros(d["shape"], dtype=d["dtype"])  # type:ignore[no-matching-overload] # loading from unsanitized data
         del d["dtype"]
         del d["shape"]
 

src/parcels/convert.py

@@ -13,6 +13,7 @@
 from __future__ import annotations
 
 import typing
+from typing import cast
 
 import numpy as np
 import xarray as xr
@@ -565,7 +566,7 @@ def _detect_vertical_coordinates(
     ValueError
         If vertical coordinates cannot be detected.
     """
-    ds_dims = set(ds.dims)
+    ds_dims = cast(set[str], set(ds.dims))
 
     # Strategy 1: Use known mappings if provided and dimensions exist
     if known_mappings is not None: