Add types sgrid

.gitattributes

@@ -1,2 +1,5 @@
 # GitHub syntax highlighting
 pixi.lock linguist-language=YAML
+*.py text eol=lf
+*.yaml text eol=lf
+*.yml text eol=lf

.pre-commit-config.yaml

@@ -10,6 +10,14 @@ repos:
     - id: check-added-large-files
       exclude_types: [yaml]
 
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v2.1.0
+  hooks:
+  - id: mypy
+    additional_dependencies:
+      - pandas-stubs
+    exclude: 'docs/source/conf\.py'
+
 - repo: https://github.com/keewis/blackdoc
   rev: v0.4.6
   hooks:

xarray_subset_grid/accessor.py

@@ -1,4 +1,5 @@
 import warnings
+from collections.abc import Iterable
 
 # from typing import Optional, Union
 import numpy as np
@@ -15,7 +16,7 @@
     UGrid,
 )
 
-_grid_impls = [
+_grid_impls: list[type[Grid]] = [
     FVCOMGrid,
     SELFEGrid,
     UGrid,
@@ -25,7 +26,7 @@
 ]
 
 
-def register_grid_impl(grid_impl: Grid, priority: int = 0):
+def register_grid_impl(grid_impl: type[Grid], priority: int = 0):
     """Register a new grid implementation.
 
     :param grid_impl: The grid implementation to register
@@ -89,7 +90,7 @@ def coords(self) -> set[str]:
         return set()
 
     @property
-    def grid_vars(self) -> set[str]:
+    def grid_vars(self) -> Iterable[str]:
         """List of grid variables.
 
         These variables are used to define the grid and thus should be

xarray_subset_grid/grid.py

@@ -26,7 +26,7 @@ def name(self) -> str:
         return "grid"
 
     @abstractmethod
-    def grid_vars(self, ds: xr.Dataset) -> list[str]:
+    def grid_vars(self, ds: xr.Dataset) -> Iterable[str]:
         """List of grid variables.
 
         These variables are used to define the grid and thus should be

xarray_subset_grid/grids/regular_grid.py

@@ -158,7 +158,8 @@ def data_vars(self, ds: xr.Dataset) -> set[str]:
     def compute_polygon_subset_selector(
         self,
         ds: xr.Dataset,
-        polygon: list[tuple[float, float]],
+        polygon: list[tuple[float, float]] | np.ndarray,
+        name: str | None = None,
     ) -> Selector:
 
         polygon = np.asarray(polygon)
@@ -173,6 +174,7 @@ def compute_bbox_subset_selector(
         self,
         ds: xr.Dataset,
         bbox: tuple[float, float, float, float],
+        name: str | None = None,
     ) -> Selector:
         bbox = normalize_bbox_x_coords(ds.cf["longitude"].values, bbox)
         selector = RegularGridBBoxSelector(bbox)

xarray_subset_grid/grids/regular_grid_2d.py

@@ -88,7 +88,7 @@ def data_vars(self, ds: xr.Dataset) -> set[str]:
         return data_vars
 
     def compute_polygon_subset_selector(
-        self, ds: xr.Dataset, polygon: list[tuple[float, float]], name: str = None
+        self, ds: xr.Dataset, polygon: list[tuple[float, float]], name: str | None = None
     ) -> Selector:
         lat = ds.cf["latitude"]
         lon = ds.cf["longitude"]

xarray_subset_grid/grids/sgrid.py

@@ -92,19 +92,15 @@ def data_vars(self, ds: xr.Dataset) -> set[str]:
         """
         grid_topology_key = ds.cf.cf_roles["grid_topology"][0]
         grid_topology = ds[grid_topology_key]
-        dims = []
-        for dims, _coords in _get_sgrid_dim_coord_names(grid_topology):
-            dims.extend(dims)
-        dims = set(dims)
+        dims = set([d for d, _ in _get_sgrid_dim_coord_names(grid_topology)])
 
         return {var for var in ds.data_vars if not set(ds[var].dims).isdisjoint(dims)}
 
     def compute_polygon_subset_selector(
-        self, ds: xr.Dataset, polygon: list[tuple[float, float]], name: str = None
+        self, ds: xr.Dataset, polygon: list[tuple[float, float]], name: str | None = None
     ) -> Selector:
         grid_topology_key = ds.cf.cf_roles["grid_topology"][0]
         grid_topology = ds[grid_topology_key]
-        dims = _get_sgrid_dim_coord_names(grid_topology)
         subset_masks: list[tuple[list[str], xr.DataArray]] = []
 
         node_info = _get_location_info_from_topology(grid_topology, "node")
@@ -114,8 +110,6 @@ def compute_polygon_subset_selector(
         unique_dims = set(node_dims)
         node_vars = [k for k in ds.variables if unique_dims.issubset(set(ds[k].dims))]
 
-        node_lon: xr.DataArray | None = None
-        node_lat: xr.DataArray | None = None
         for c in node_coords:
             if "lon" in ds[c].standard_name.lower():
                 node_lon = ds[c]
@@ -144,14 +138,14 @@ def compute_polygon_subset_selector(
             info = _get_location_info_from_topology(grid_topology, s)
             dims = info["dims"]
             coords = info["coords"]
-            unique_dims = set(dims)
-            vars = [k for k in ds.variables if unique_dims.issubset(set(ds[k].dims))]
+            _unique_dims = set(dims)
+            vars = [k for k in ds.variables if _unique_dims.issubset(set(ds[k].dims))]
 
-            lon: xr.DataArray | None = None
             for c in coords:
                 if "lon" in ds[c].standard_name.lower():
                     lon = ds[c]
             padding = info["padding"]
+            arranged_padding: list[int | str] = []
             arranged_padding = [padding[d] for d in lon.dims]
             arranged_padding = [0 if p == "none" or p == "low" else 1 for p in arranged_padding]
             mask = np.zeros(lon.shape, dtype=bool)
@@ -173,17 +167,16 @@ def compute_polygon_subset_selector(
 
 def _get_location_info_from_topology(grid_topology: xr.DataArray, location) -> dict[str, str]:
     """Get the dimensions and coordinates for a given location from the grid_topology"""
-    rdict = {}
     dim_str = grid_topology.attrs.get(f"{location}_dimensions", None)
     coord_str = grid_topology.attrs.get(f"{location}_coordinates", None)
     if dim_str is None or coord_str is None:
         raise ValueError(f"Could not find {location} dimensions or coordinates")
     # Remove padding for now
-    dims_only = " ".join([v for v in dim_str.split(" ") if "(" not in v and ")" not in v])
-    if ":" in dims_only:
-        dims_only = [s.replace(":", "") for s in dims_only.split(" ") if ":" in s]
+    _dims_only = " ".join([v for v in dim_str.split(" ") if "(" not in v and ")" not in v])
+    if ":" in _dims_only:
+        dims_only = [s.replace(":", "") for s in _dims_only.split(" ") if ":" in s]
     else:
-        dims_only = dims_only.split(" ")
+        dims_only = _dims_only.split(" ")
 
     padding = dim_str.replace(":", "").split(")")
     pdict = {}
@@ -201,10 +194,11 @@ def _get_location_info_from_topology(grid_topology: xr.DataArray, location) -> d
         pdict[dims_only[0]] = "none"
         pdict[dims_only[1]] = "none"
 
-    rdict["dims"] = dims_only
-    rdict["coords"] = coord_str.split(" ")
-    rdict["padding"] = pdict
-    return rdict
+    return {
+        "dims": dims_only,
+        "coords": coord_str.split(" "),
+        "padding": pdict,
+    }
 
 
 def _get_sgrid_dim_coord_names(
@@ -220,11 +214,11 @@ def _get_sgrid_dim_coord_names(
     for k, v in grid_topology.attrs.items():
         if "_dimensions" in k:
             # Remove padding for now
-            d = " ".join([v for v in v.split(" ") if "(" not in v and ")" not in v])
-            if ":" in d:
-                d = [d.replace(":", "") for d in d.split(" ") if ":" in d]
+            _d = " ".join([v for v in v.split(" ") if "(" not in v and ")" not in v])
+            if ":" in _d:
+                d = [_d.replace(":", "") for _d in _d.split(" ") if ":" in _d]
             else:
-                d = d.split(" ")
+                d = _d.split(" ")
             dims.append(d)
         elif "_coordinates" in k:
             coords.append(v.split(" "))

xarray_subset_grid/grids/ugrid.py

@@ -171,7 +171,7 @@ def data_vars(self, ds: xr.Dataset) -> set[str]:
         return data_vars
 
     def compute_polygon_subset_selector(
-        self, ds: xr.Dataset, polygon: list[tuple[float, float]], name: str = None
+        self, ds: xr.Dataset, polygon: list[tuple[float, float]], name: str | None = None
     ) -> Selector:
         # For this grid type, we find all nodes that are connected to elements that are inside
         # the polygon. To do this, we first find all nodes that are inside the polygon and then