Types take 2

.pre-commit-config.yaml

@@ -16,14 +16,20 @@ repos:
   - id: mypy
     # In the next typing PR we will add the extra dependencies for a better assessment, then
     # we'll try strict, then removing ignore missing imports.
-    # args: [--strict, --ignore-missing-imports]
+    args: [ --install-types, --non-interactive ]
     additional_dependencies:
-      - pandas-stubs
-    #   - cf-xarray
-    #   - netcdf4
-    #   - xarray>=2024.6
-
-    exclude: 'docs/source/conf\.py'
+      - cf-xarray
+      - netcdf4
+      - xarray>=2024.6
+      - matplotlib
+    exclude: |
+      (?x)(
+        ^build/
+        | conf\.py$
+        | ^tests/
+        | ^docs/
+        | ^profiling/
+      )
 
 - repo: https://github.com/keewis/blackdoc
   rev: v0.4.6

xarray_subset_grid/accessor.py

@@ -1,10 +1,13 @@
 import warnings
 from collections.abc import Iterable
+from typing import TYPE_CHECKING
 
-# from typing import Optional, Union
 import numpy as np
 import xarray as xr
 
+if TYPE_CHECKING:
+    from xarray.core.coordinates import DatasetCoordinates
+
 from xarray_subset_grid.grid import Grid
 from xarray_subset_grid.grids import (
     FVCOMGrid,
@@ -84,7 +87,7 @@ def data_vars(self) -> set[str]:
         return set()
 
     @property
-    def coords(self) -> set[str]:
+    def coords(self) -> set[str] | "DatasetCoordinates":
         if self._ds:
             return self._ds.coords
         return set()

xarray_subset_grid/grids/fvcom_grid.py

@@ -96,11 +96,11 @@ def subset_vertical_level(
         selections = {}
         if "siglay" in ds.dims:
             siglay = ds["siglay"].isel(node=0)
-            elevation_index = int(np.absolute(siglay - level).argmin(axis=0).values)
+            elevation_index = int(np.absolute((siglay - level).to_numpy()).argmin(axis=0))
             selections["siglay"] = elevation_index
         if "siglev" in ds.dims:
             siglev = ds["siglev"].isel(node=0)
-            elevation_index = int(np.absolute(siglev - level).argmin(axis=0).values)
+            elevation_index = int(np.absolute((siglev - level).to_numpy()).argmin(axis=0))
             selections["siglev"] = elevation_index
 
         return ds.isel(selections)
@@ -127,13 +127,13 @@ def subset_vertical_levels(
         if "siglay" in ds.dims:
             siglay = ds["siglay"].isel(node=0)
             elevation_indexes = [
-                int(np.absolute(siglay - level).argmin(axis=0).values) for level in levels
+                int(np.absolute((siglay - level).to_numpy()).argmin(axis=0)) for level in levels
             ]
             selections["siglay"] = slice(elevation_indexes[0], elevation_indexes[1])
         if "siglev" in ds.dims:
             siglev = ds["siglev"].isel(node=0)
             elevation_index = [
-                int(np.absolute(siglev - level).argmin(axis=0).values) for level in levels
+                int(np.absolute((siglev - level).to_numpy()).argmin(axis=0)) for level in levels
             ]
             selections["siglev"] = slice(elevation_index[0], elevation_index[1])
 

xarray_subset_grid/grids/sgrid.py

@@ -44,10 +44,9 @@ def select(self, ds: xr.Dataset) -> xr.Dataset:
             ds_out.append(ds_subset)
 
         # Merge the subsetted datasets
-        ds_out = xr.merge(ds_out)
+        _ds_out = xr.merge(ds_out)
 
-        ds_out = ds_out.assign({self._grid_topology_key: self._grid_topology})
-        return ds_out
+        return _ds_out.assign({self._grid_topology_key: self._grid_topology})
 
 
 class SGrid(Grid):

xarray_subset_grid/grids/ugrid.py

@@ -477,12 +477,12 @@ def assign_ugrid_topology(
         raise ValueError(f"start_index must be 0 or 1, not {mesh.start_index}")
 
     # assign the start_index to all the grid variables
-    for var in (v for v in ALL_MESH_VARS if "connectivity" in v):
-        var_name = getattr(mesh, var)
+    for _var in (v for v in ALL_MESH_VARS if "connectivity" in v):
+        var_name = getattr(mesh, _var)
         if var_name:
             try:
-                var = ds[var_name]
-                var.attrs.setdefault("start_index", mesh.start_index)
+                ds_var = ds[var_name]
+                ds_var.attrs.setdefault("start_index", mesh.start_index)
             except KeyError:
                 warnings.warn(f"{var_name} in mesh_topology variable, but not in dataset")
 

xarray_subset_grid/utils.py

@@ -1,7 +1,9 @@
 from datetime import datetime
 
 import cf_xarray  # noqa
-import cftime
+
+# See https://github.com/Unidata/cftime/issues/349
+import cftime  # type: ignore[import-untyped]
 import numpy as np
 import xarray as xr
 from dateutil.parser import parse as parsetime