Merge branch 'main' into better-community

Author: VeckoTheGecko

.github/workflows/ci.yml

@@ -170,8 +170,6 @@ jobs:
     name: "TypeChecking: pixi run typing"
     runs-on: ubuntu-latest
     needs: [cache-pixi-lock]
-    # TODO v4: Enable typechecking again
-    if: false
     steps:
       - name: Checkout
         uses: actions/checkout@v5
@@ -185,14 +183,8 @@ jobs:
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
       - name: Typechecking
-        run: |
-          pixi run typing --non-interactive --html-report mypy-report
-      - name: Upload test results
-        if: ${{ always() }} # Upload even on mypy error
-        uses: actions/upload-artifact@v7
-        with:
-          name: Mypy report
-          path: mypy-report
+        run: | # TODO: Remove `|| true` once typechecking is stable
+          pixi run typing --output-format github || true
   build-and-upload-nightly-parcels: # for alpha testing
     needs: [cache-pixi-lock]
     permissions:
@@ -220,7 +212,7 @@ jobs:
         run: |
           for pkg in $(find output -type f \( -name "*.conda" -o -name "*.tar.bz2" \) ); do
             echo "Uploading ${pkg}"
-            rattler-build upload prefix -c parcels "${pkg}"
+            pixi run -e rattler-build rattler-build upload prefix -c parcels "${pkg}"
           done
         env:
           PREFIX_API_KEY: ${{ secrets.PREFIX_API_KEY }}

docs/development/index.md

@@ -132,8 +132,8 @@ See below for more Pixi commands relevant to development.
 
 **Code quality**
 
-- `pixi run lint` - Run pre-commit hooks on all files (includes formatting, linting, and other code quality checks)
-- `pixi run typing` - Run mypy type checking on the codebase
+- `pixi run lint` - Run [pre-commit](https://pre-commit.com/) hooks on all files (includes formatting, linting, and other code quality checks)
+- `pixi run typing` - Run [ty](https://docs.astral.sh/ty/) type checking on the codebase
 
 **Different environments**
 

pixi.toml

@@ -82,6 +82,25 @@ trajan = "*"
 matplotlib-base = ">=2.0.2"
 gsw = "*"
 
+[feature.devtools.dependencies]
+pdbpp = "*"
+line_profiler = "*"
+memory_profiler = "*"
+snakeviz = "*"
+icecream = "*"
+ipykernel = "*"
+snoop = "*"
+pyinstrument = "*"
+
+[feature.devtools.target.linux-64.dependencies]
+memray = "*"
+
+[feature.devtools.target.osx-64.dependencies]
+memray = "*"
+
+[feature.devtools.target.osx-arm64.dependencies]
+memray = "*"
+
 [feature.docs.dependencies]
 parcels = { path = "." }
 numpydoc = "*"
@@ -113,12 +132,10 @@ numpydoc = "*"
 numpydoc-lint = { cmd = "python tools/numpydoc-public-api.py", description = "Lint public API docstrings with numpydoc." }
 
 [feature.typing.dependencies]
-mypy = "*"
-lxml = "*" # in CI
-types-tqdm = "*"
+ty = "*"
 
 [feature.typing.tasks]
-typing = { cmd = "mypy src/parcels --install-types", description = "Run static type checking with mypy." }
+typing = { cmd = "ty check", description = "Run static type checking with ty." }
 
 
 [environments]
@@ -128,6 +145,7 @@ default = { features = [
   "typing",
   "pre-commit",
   "numpydoc",
+  "devtools",
 ], solve-group = "main" }
 test = { features = ["test"], solve-group = "main" }
 test-minimum = { features = ["test", "minimum"] }

pyproject.toml

@@ -157,25 +157,8 @@ convention = "numpy"
 [tool.ruff.lint.isort]
 known-first-party = ["parcels"]
 
-[tool.mypy]
-files = [
-  "parcels/_typing.py",
-  "parcels/tools/*.py",
-  "parcels/grid.py",
-  "parcels/field.py",
-  "parcels/fieldset.py",
-]
-
-[[tool.mypy.overrides]]
-module = [
-  "parcels._version_setup",
-  "mpi4py",
-  "scipy.spatial",
-  "sklearn.cluster",
-  "zarr",
-  "cftime",
-  "pykdtree.kdtree",
-  "netCDF4",
-  "pooch",
+[tool.ty.src]
+include = ["./src/"]
+exclude = [
+  "./src/parcels/interpolators/", # ignore for now
 ]
-ignore_missing_imports = true

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]