Fix type annotations

Author: VeckoTheGecko

src/parcels/_core/fieldset.py

@@ -284,14 +284,17 @@ def from_fesom2(ds: ux.UxDataset):
 
         fields: dict[str, Field | VectorField] = {}
         if "U" in ds.data_vars and "V" in ds.data_vars:
-            fields["U"] = Field("U", ds["U"], grid, _select_uxinterpolator(ds["U"]))
-            fields["V"] = Field("V", ds["V"], grid, _select_uxinterpolator(ds["U"]))
+            field_U = Field("U", ds["U"], grid, _select_uxinterpolator(ds["U"]))
+            field_V = Field("V", ds["V"], grid, _select_uxinterpolator(ds["U"]))
+            fields["U"] = field_U
+            fields["V"] = field_V
 
             if "W" in ds.data_vars:
-                fields["W"] = Field("W", ds["W"], grid, _select_uxinterpolator(ds["U"]))
-                fields["UVW"] = VectorField("UVW", fields["U"], fields["V"], fields["W"])
+                field_W = Field("W", ds["W"], grid, _select_uxinterpolator(ds["U"]))
+                fields["W"] = field_W
+                fields["UVW"] = VectorField("UVW", field_U, field_V, field_W)
             else:
-                fields["UV"] = VectorField("UV", fields["U"], fields["V"])
+                fields["UV"] = VectorField("UV", field_U, field_V)
 
         for varname in set(ds.data_vars) - set(fields.keys()):
             fields[varname] = Field(varname, ds[varname], grid, _select_uxinterpolator(ds[varname]))

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,14 +62,14 @@ 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: float):
     """Find and return the index and relative coordinate in the time array associated with a given time.
 
     Parameters
     ----------
     field: Field
 
-    time: datetime
+    time: float
         This is the amount of time, in seconds (time_delta), in unix epoch
     Note that we normalize to either the first or the last index
     if the sampled value is outside the time value range.
@@ -172,6 +171,8 @@ def _search_indices_curvilinear_2d(
     """
     if np.any(xi):
         # If an initial guess is provided, we first perform a point in cell check for all guessed indices
+        assert xi is not None
+        assert yi is not None
         is_in_cell, coords = curvilinear_point_in_cell(grid, y, x, yi, xi)
         y_check = y[is_in_cell == 0]
         x_check = x[is_in_cell == 0]

src/parcels/_core/xgrid.py

@@ -380,7 +380,7 @@ def get_axis_dim_mapping(self, dims: list[str]) -> dict[_XGRID_AXES, str]:
         return result
 
 
-def get_axis_from_dim_name(axes: _XGCM_AXES, dim: str) -> _XGCM_AXIS_DIRECTION | None:
+def get_axis_from_dim_name(axes: _XGCM_AXES, dim: Hashable) -> _XGCM_AXIS_DIRECTION | None:
     """For a given dimension name in a grid, returns the direction axis it is on."""
     for axis_name, axis in axes.items():
         if dim in axis.coords.values():