improve types hints

src/xarray_grass/__init__.py

@@ -1,10 +1,10 @@
-from xarray_grass.grass_interface import GrassConfig as GrassConfig
-from xarray_grass.grass_interface import GrassInterface as GrassInterface
-from xarray_grass.xarray_grass import GrassBackendEntrypoint as GrassBackendEntrypoint
+from xarray_grass.coord_utils import RegionData as RegionData
 from xarray_grass.grass_backend_array import (
     GrassSTDSBackendArray as GrassSTDSBackendArray,
 )
+from xarray_grass.grass_interface import GrassConfig as GrassConfig
+from xarray_grass.grass_interface import GrassInterface as GrassInterface
 from xarray_grass.to_grass import to_grass as to_grass
-from xarray_grass.coord_utils import RegionData as RegionData
+from xarray_grass.xarray_grass import GrassBackendEntrypoint as GrassBackendEntrypoint
 
 __version__ = "0.4.0"

src/xarray_grass/coord_utils.py

@@ -13,12 +13,12 @@
 GNU General Public License for more details.
 """
 
-from collections import namedtuple
-from typing import Mapping
+from collections.abc import Mapping
+from typing import NamedTuple
+
 import numpy as np
 import xarray as xr  # For type hinting xr.DataArray
 
-
 region_type_dict = {
     "projection": str,
     "zone": str,
@@ -41,14 +41,34 @@
     "cells": int,
     "cells3": int,
 }
-RegionData = namedtuple(
-    "RegionData",
-    region_type_dict.keys(),
-    defaults=[None for _ in region_type_dict.keys()],
-)
 
 
-def get_region_from_xarray(data_array: xr.DataArray, dims: Mapping[str, str]) -> dict:
+class RegionData(NamedTuple):
+    projection: str | None = None
+    zone: str | None = None
+    n: float | None = None
+    s: float | None = None
+    w: float | None = None
+    e: float | None = None
+    t: float | None = None
+    b: float | None = None
+    nsres: float | None = None
+    nsres3: float | None = None
+    ewres: float | None = None
+    ewres3: float | None = None
+    tbres: float | None = None
+    rows: int | None = None
+    rows3: int | None = None
+    cols: int | None = None
+    cols3: int | None = None
+    depths: int | None = None
+    cells: int | None = None
+    cells3: int | None = None
+
+
+def get_region_from_xarray(
+    data_array: xr.DataArray, dims: Mapping[str, str]
+) -> RegionData:
     """
     Calculates GRASS GIS region parameters from an xarray DataArray.
 
@@ -73,8 +93,8 @@ def get_region_from_xarray(data_array: xr.DataArray, dims: Mapping[str, str]) ->
 
     Returns
     -------
-    dict
-        A dictionary containing GRASS region parameters:
+    RegionData
+        GRASS region parameters:
         'n', 's', 'e', 'w': float or None (geographical limits)
         't', 'b': float or None (top, bottom limits for 3D)
         'nsres', 'ewres': float or None (2D resolutions)
@@ -84,7 +104,7 @@ def get_region_from_xarray(data_array: xr.DataArray, dims: Mapping[str, str]) ->
     """
     region = {}
 
-    def _calculate_res(coords_arr_np: np.ndarray) -> float | None:
+    def _calculate_res(coords_arr_np: np.ndarray | None) -> float | None:
         if coords_arr_np is not None and len(coords_arr_np) >= 2:
             # Ensure consistent dtype for subtraction, then convert to float
             res = np.abs(
@@ -95,7 +115,7 @@ def _calculate_res(coords_arr_np: np.ndarray) -> float | None:
 
     # Determine if it's 3D based on presence of z-coordinate name in dims and data_array
     z_name = dims.get("z")
-    is_3d = z_name and z_name in data_array.coords
+    is_3d = z_name is not None and z_name in data_array.coords
 
     x_coords_np, y_coords_np, z_coords_np = None, None, None
 
@@ -177,4 +197,16 @@ def _calculate_res(coords_arr_np: np.ndarray) -> float | None:
         region["b"] = float(z_coords_np[0] - region["tbres"] / 2)
         region["t"] = float(z_coords_np[-1] + region["tbres"] / 2)
 
-    return RegionData(**region)
+    return RegionData(
+        n=region.get("n"),
+        s=region.get("s"),
+        w=region.get("w"),
+        e=region.get("e"),
+        t=region.get("t"),
+        b=region.get("b"),
+        nsres=region.get("nsres"),
+        nsres3=region.get("nsres3"),
+        ewres=region.get("ewres"),
+        ewres3=region.get("ewres3"),
+        tbres=region.get("tbres"),
+    )

src/xarray_grass/grass_backend_array.py

@@ -13,47 +13,47 @@
 """
 
 from __future__ import annotations
-from typing import TYPE_CHECKING
+
 import threading
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
-import xarray as xr
-
 from xarray.backends import BackendArray
+from xarray.core import indexing
 
 if TYPE_CHECKING:
-    from xarray_grass.grass_interface import GrassInterface
+    from xarray_grass.grass_interface import GrassInterface, MapData
 
 
 class GrassSTDSBackendArray(BackendArray):
     """Lazy loading of grass Space-Time DataSets (multiple maps in time series)"""
 
     def __init__(
         self,
-        shape,
-        dtype,
-        map_list: list,  # List of map metadata objects
-        map_type: str,
+        shape: tuple[int, ...],
+        dtype: np.dtype[Any],
+        map_list: list[MapData],
+        map_type: Literal["raster", "raster3d"],
         grass_interface: GrassInterface,
-    ):
+    ) -> None:
         self.shape = shape
         self.dtype = dtype
         self._lock = threading.Lock()
-        self.map_list = map_list  # List with .id attribute
-        self.map_type = map_type  # "raster" or "raster3d"
+        self.map_list = map_list
+        self.map_type = map_type
         self.grass_interface = grass_interface
-        self._cached_maps = {}  # Cache loaded maps by index
+        self._cached_maps: dict[int, np.ndarray] = {}
 
-    def __getitem__(self, key: xr.core.indexing.ExplicitIndexer) -> np.typing.ArrayLike:
+    def __getitem__(self, key: indexing.ExplicitIndexer) -> np.ndarray:
         """takes in input an index and returns a NumPy array"""
-        return xr.core.indexing.explicit_indexing_adapter(
+        return indexing.explicit_indexing_adapter(
             key,
             self.shape,
-            xr.core.indexing.IndexingSupport.BASIC,
+            indexing.IndexingSupport.BASIC,
             self._raw_indexing_method,
         )
 
-    def _raw_indexing_method(self, key: tuple):
+    def _raw_indexing_method(self, key: tuple[Any, ...]) -> np.ndarray:
         """Load only the maps needed for the requested slice"""
         with self._lock:
             # key is a tuple of slices/indices for each dimension
@@ -70,7 +70,7 @@ def _raw_indexing_method(self, key: tuple):
                 time_indices = list(time_key)
 
             # Load only the needed maps
-            result_list = []
+            result_list: list[np.ndarray] = []
             for t_idx in time_indices:
                 if t_idx not in self._cached_maps:
                     map_data = self.map_list[t_idx]