Fix typechecking errors

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
@@ -44,7 +43,7 @@
     """
     # TODO v4: We probably rework this to deal with 0D arrays before this point (as we already know field dimensionality)
     if len(arr) < 2:
         return np.zeros(shape=x.shape, dtype=np.int32), np.zeros_like(x)
     index = np.clip(np.searchsorted(arr, x, side="right") - 1, 0, len(arr) - 2)
     # Use broadcasting to avoid repeated array access
     arr_index = arr[index]
@@ -63,7 +62,7 @@
     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

@@ -376,7 +376,7 @@
         for dim in dims:
             axis = get_axis_from_dim_name(self.xgcm_grid.axes, dim)
             if axis in self.axes:  # Only include spatial axes (X, Y, Z)
                 result[cast(ptyping.XgridAxis, axis)] = dim
         return result
 
 
@@ -394,7 +394,7 @@
         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: