fixing mypy errors

Author: VeckoTheGecko

parcels/_core/utils/time.py

@@ -6,10 +6,10 @@
 import cftime
 import numpy as np
 
-T = TypeVar("T", datetime, cftime.datetime)
-
 if TYPE_CHECKING:
-    from parcels._typing import DatetimeLike
+    from parcels._typing import TimeLike
+
+T = TypeVar("T", bound="TimeLike")
 
 
 class TimeInterval:
@@ -92,7 +92,7 @@ def is_compatible(
 
 
 def get_datetime_type_calendar(
-    example_datetime: DatetimeLike,
+    example_datetime: TimeLike,
 ) -> tuple[type, str | None]:
     """Get the type and calendar of a datetime object.
 

parcels/_core/utils/unstructured.py

@@ -1,12 +1,10 @@
-from collections.abc import Hashable
-
 DIM_TO_VERTICAL_LOCATION_MAP = {
     "nz1": "center",
     "nz": "face",
 }
 
 
-def get_vertical_location_from_dims(dims: tuple[Hashable, ...]):
+def get_vertical_location_from_dims(dims: tuple[str, ...]):
     """
     Determine the vertical location of the field based on the uxarray.UxDataArray object variables.
 

parcels/_index_search.py

@@ -356,7 +356,7 @@ def _search_indices_curvilinear_2d(
 
 
 ## TODO :  Still need to implement the search_indices_curvilinear
-def _search_indices_curvilinear(field: Field, time, z, y, x, ti, particle=None, search2D=False):
+def _search_indices_curvilinear(field, time, z, y, x, ti, particle=None, search2D=False):
     if particle:
         zi, yi, xi = field.unravel_index(particle.ei)
     else:

parcels/_typing.py

@@ -34,7 +34,7 @@
 VectorType = Literal["3D", "3DSigma", "2D"] | None  # corresponds with `vector_type`
 GridIndexingType = Literal["pop", "mom5", "mitgcm", "nemo", "croco"]  # corresponds with `gridindexingtype`
 NetcdfEngine = Literal["netcdf4", "xarray", "scipy"]
-DatetimeLike = datetime | cftime_datetime | np.datetime64
+TimeLike = datetime | cftime_datetime | np.datetime64
 
 KernelFunction = Callable[..., None]
 

parcels/fieldset.py

@@ -15,7 +15,7 @@
 from parcels.xgrid import XGrid
 
 if TYPE_CHECKING:
-    from parcels._typing import DatetimeLike
+    from parcels._typing import TimeLike
     from parcels.basegrid import BaseGrid
 __all__ = ["FieldSet"]
 
@@ -57,7 +57,7 @@ def __init__(self, fields: list[Field | VectorField]):
         assert_compatible_calendars(fields)
 
         self.fields = {f.name: f for f in fields}
-        self.constants = {}
+        self.constants: dict[str, float] = {}
 
     def __getattr__(self, name):
         """Get the field by name. If the field is not found, check if it's a constant."""
@@ -236,7 +236,7 @@ class CalendarError(Exception):  # TODO: Move to a parcels errors module
     """Exception raised when the calendar of a field is not compatible with the rest of the Fields. The user should ensure that they only add fields to a FieldSet that have compatible CFtime calendars."""
 
 
-def assert_compatible_calendars(fields: Iterable[Field]):
+def assert_compatible_calendars(fields: Iterable[Field | VectorField]):
     time_intervals = [f.time_interval for f in fields if f.time_interval is not None]
 
     if len(time_intervals) == 0:  # All time intervals are none
@@ -253,13 +253,13 @@ def assert_compatible_calendars(fields: Iterable[Field]):
             raise CalendarError(msg)
 
 
-def _datetime_to_msg(example_datetime: DatetimeLike) -> str:
+def _datetime_to_msg(example_datetime: TimeLike) -> str:
     datetime_type, calendar = get_datetime_type_calendar(example_datetime)
     msg = str(datetime_type)
     if calendar is not None:
         msg += f" with cftime calendar {calendar}'"
     return msg
 
 
-def _format_calendar_error_message(field: Field, reference_datetime: DatetimeLike) -> str:
+def _format_calendar_error_message(field: Field, reference_datetime: TimeLike) -> str:
     return f"Expected field {field.name!r} to have calendar compatible with datetime object {_datetime_to_msg(reference_datetime)}. Got field with calendar {_datetime_to_msg(field.time_interval.left)}. Have you considered using xarray to update the time dimension of the dataset to have a compatible calendar?"