Improve typechecking (#2541)

Co-authored-by: Claude <[email protected]>

Author: VeckoTheGecko

pyproject.toml

@@ -159,3 +159,6 @@ known-first-party = ["parcels"]
 
 [tool.ty.src]
 include = ["./src/"]
+exclude = [
+  "./src/parcels/interpolators/", # ignore for now
+]

src/parcels/_compat.py

@@ -8,13 +8,13 @@
 KMeans: Any | None = None
 
 try:
-    from mpi4py import MPI  # type: ignore[no-redef]
+    from mpi4py import MPI  # type: ignore[import-untyped,no-redef]
 except ModuleNotFoundError:
     pass
 
 # KMeans is used in MPI. sklearn not installed by default
 try:
-    from sklearn.cluster import KMeans  # type: ignore[no-redef]
+    from sklearn.cluster import KMeans
 except ModuleNotFoundError:
     pass
 

src/parcels/_core/field.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import warnings
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from datetime import datetime
 
 import numpy as np
@@ -428,7 +428,7 @@ def _assert_valid_uxdataarray(data: ux.UxDataArray):
             )
 
 
-def _assert_compatible_combination(data: xr.DataArray | ux.UxDataArray, grid: ux.Grid | XGrid):
+def _assert_compatible_combination(data: xr.DataArray | ux.UxDataArray, grid: UxGrid | XGrid):
     if isinstance(data, ux.UxDataArray):
         if not isinstance(grid, UxGrid):
             raise ValueError(
@@ -448,7 +448,7 @@ def _get_time_interval(data: xr.DataArray | ux.UxDataArray) -> TimeInterval | No
     return TimeInterval(data.time.values[0], data.time.values[-1])
 
 
-def _assert_same_time_interval(fields: list[Field]) -> None:
+def _assert_same_time_interval(fields: Sequence[Field]) -> None:
     if len(fields) == 0:
         return
 

src/parcels/_core/fieldset.py

@@ -223,7 +223,7 @@ def from_ugrid_conventions(cls, ds: ux.UxDataset, mesh: str = "spherical"):
                 )
 
         for varname in set(ds.data_vars) - set(fields.keys()):
-            fields[varname] = Field(varname, ds[varname], grid, _select_uxinterpolator(ds[varname]))
+            fields[varname] = Field(str(varname), ds[varname], grid, _select_uxinterpolator(ds[varname]))
 
         return cls(list(fields.values()))
 
@@ -319,7 +319,7 @@ def from_sgrid_conventions(
                 )
 
         for varname in set(ds.data_vars) - set(fields.keys()) - skip_vars:
-            fields[varname] = Field(varname, ds[varname], grid, XLinear)
+            fields[varname] = Field(str(varname), ds[varname], grid, XLinear)
 
         return cls(list(fields.values()))
 
@@ -353,7 +353,7 @@ def _datetime_to_msg(example_datetime: TimeLike) -> str:
     return msg
 
 
-def _format_calendar_error_message(field: Field, reference_datetime: TimeLike) -> str:
+def _format_calendar_error_message(field: Field | VectorField, 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?"
 
 

src/parcels/_core/index_search.py

@@ -10,7 +10,7 @@
 
 if TYPE_CHECKING:
     from parcels._core.field import Field
-    from parcels.xgrid import XGrid
+    from parcels._core.xgrid import XGrid
 
 
 GRID_SEARCH_ERROR = -3
@@ -19,7 +19,7 @@
 
 
 def _search_1d_array(
-    arr: np.array,
+    arr: np.ndarray,
     x: float,
 ) -> tuple[int, int]:
     """

src/parcels/_core/particle.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import operator
-from typing import Literal
+from typing import Any, Literal
 
 import numpy as np
 
@@ -37,7 +37,7 @@ class Variable:
     def __init__(
         self,
         name,
-        dtype: np.dtype = np.float32,
+        dtype: np.dtype[Any] | type[np.generic] = np.float32,
         initial=0,
         to_write: bool | Literal["once"] = True,
         attrs: dict | None = None,
@@ -122,7 +122,7 @@ def _assert_no_duplicate_variable_names(*, existing_vars: list[Variable], new_va
             raise ValueError(f"Variable name already exists: {var.name}")
 
 
-def get_default_particle(spatial_dtype: np.float32 | np.float64) -> ParticleClass:
+def get_default_particle(spatial_dtype: type[np.float32] | type[np.float64]) -> ParticleClass:
     if spatial_dtype not in [np.float32, np.float64]:
         raise ValueError(f"spatial_dtype must be np.float32 or np.float64. Got {spatial_dtype=!r}")
 
@@ -177,7 +177,7 @@ def create_particle_data(
     nparticles: int,
     ngrids: int,
     time_interval: TimeInterval,
-    initial: dict[str, np.array] | None = None,
+    initial: dict[str, np.ndarray] | None = None,
 ):
     if initial is None:
         initial = {}

src/parcels/_core/particlefile.py

@@ -196,10 +196,10 @@ def _write_particle_data(self, *, particle_data, pclass, time_interval, time, in
         if self.create_new_zarrfile:
             if self.chunks is None:
                 self._chunks = (nparticles, 1)
-            if (self._maxids > len(ids)) or (self._maxids > self.chunks[0]):  # type: ignore[index]
-                arrsize = (self._maxids, self.chunks[1])  # type: ignore[index]
+            if (self._maxids > len(ids)) or (self._maxids > self.chunks[0]):
+                arrsize = (self._maxids, self.chunks[1])
             else:
-                arrsize = (len(ids), self.chunks[1])  # type: ignore[index]
+                arrsize = (len(ids), self.chunks[1])
             ds = xr.Dataset(
                 attrs=self.metadata,
                 coords={"trajectory": ("trajectory", pids), "obs": ("obs", np.arange(arrsize[1], dtype=np.int32))},
@@ -221,7 +221,7 @@ def _write_particle_data(self, *, particle_data, pclass, time_interval, time, in
                         data[ids, 0] = particle_data[var.name][indices_to_write]
                         dims = ["trajectory", "obs"]
                     ds[var.name] = xr.DataArray(data=data, dims=dims, attrs=attrs[var.name])
-                    ds[var.name].encoding["chunks"] = self.chunks[0] if var.to_write == "once" else self.chunks  # type: ignore[index]
+                    ds[var.name].encoding["chunks"] = self.chunks[0] if var.to_write == "once" else self.chunks
             ds.to_zarr(store, mode="w")
             self._create_new_zarrfile = False
         else:
@@ -234,7 +234,7 @@ def _write_particle_data(self, *, particle_data, pclass, time_interval, time, in
                     if len(once_ids) > 0:
                         Z[var.name].vindex[ids_once] = particle_data[var.name][indices_to_write_once]
                 else:
-                    if max(obs) >= Z[var.name].shape[1]:  # type: ignore[type-var]
+                    if max(obs) >= Z[var.name].shape[1]:
                         self._extend_zarr_dims(Z[var.name], store, dtype=var.dtype, axis=1)
                     Z[var.name].vindex[ids, obs] = particle_data[var.name][indices_to_write]
 

src/parcels/_core/utils/interpolation.py

@@ -5,7 +5,7 @@
 
 from parcels._typing import Mesh
 
-__all__ = []  # type: ignore
+__all__ = []
 
 
 def phi1D_lin(xsi: float) -> list[float]:

src/parcels/_core/utils/sgrid.py

@@ -15,7 +15,7 @@
 import re
 from collections.abc import Hashable, Iterable
 from dataclasses import dataclass
-from typing import Any, Literal, Protocol, Self, overload
+from typing import Any, Literal, Protocol, Self, cast, overload
 
 import xarray as xr
 
@@ -149,13 +149,13 @@ def __eq__(self, other: Any) -> bool:
         return self.to_attrs() == other.to_attrs()
 
     @classmethod
-    def from_attrs(cls, attrs):
+    def from_attrs(cls, attrs):  # type: ignore[override]
         try:
             return cls(
                 cf_role=attrs["cf_role"],
                 topology_dimension=attrs["topology_dimension"],
-                node_dimensions=load_mappings(attrs["node_dimensions"]),
-                face_dimensions=load_mappings(attrs["face_dimensions"]),
+                node_dimensions=cast(tuple[Dim, Dim], load_mappings(attrs["node_dimensions"])),
+                face_dimensions=cast(tuple[DimDimPadding, DimDimPadding], load_mappings(attrs["face_dimensions"])),
                 node_coordinates=maybe_load_mappings(attrs.get("node_coordinates")),
                 vertical_dimensions=maybe_load_mappings(attrs.get("vertical_dimensions")),
             )
@@ -176,7 +176,7 @@ def to_attrs(self) -> dict[str, str | int]:
         return d
 
     def rename(self, names_dict: dict[str, str]) -> Self:
-        return _metadata_rename(self, names_dict)
+        return cast(Self, _metadata_rename(self, names_dict))
 
     def get_value_by_id(self, id: str) -> str:
         """In the SGRID specification for 2D grids, different parts of the spec are identified by different "ID"s.
@@ -262,13 +262,15 @@ def __eq__(self, other: Any) -> bool:
         return self.to_attrs() == other.to_attrs()
 
     @classmethod
-    def from_attrs(cls, attrs):
+    def from_attrs(cls, attrs):  # type: ignore[override]
         try:
             return cls(
                 cf_role=attrs["cf_role"],
                 topology_dimension=attrs["topology_dimension"],
-                node_dimensions=load_mappings(attrs["node_dimensions"]),
-                volume_dimensions=load_mappings(attrs["volume_dimensions"]),
+                node_dimensions=cast(tuple[Dim, Dim, Dim], load_mappings(attrs["node_dimensions"])),
+                volume_dimensions=cast(
+                    tuple[DimDimPadding, DimDimPadding, DimDimPadding], load_mappings(attrs["volume_dimensions"])
+                ),
                 node_coordinates=maybe_load_mappings(attrs.get("node_coordinates")),
             )
         except Exception as e:
@@ -286,7 +288,7 @@ def to_attrs(self) -> dict[str, str | int]:
         return d
 
     def rename(self, dims_dict: dict[str, str]) -> Self:
-        return _metadata_rename(self, dims_dict)
+        return cast(Self, _metadata_rename(self, dims_dict))
 
     def get_value_by_id(self, id: str) -> str:
         """In the SGRID specification for 3D grids, different parts of the spec are identified by different "ID"s.

src/parcels/_core/uxgrid.py

@@ -18,7 +18,7 @@ class UxGrid(BaseGrid):
     for interpolation on unstructured grids.
     """
 
-    def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh) -> UxGrid:
+    def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh) -> None:
         """
         Initializes the UxGrid with a uxarray grid and vertical coordinate array.