improving type hints related to numpy

.github/workflows/type-check.yml

@@ -3,7 +3,7 @@ name: type checking
 on:
   pull_request:
     paths:
-      - yt_experiments/**/*.py
+      - yt_xarray/**/*.py
       - pyproject.toml
       - requirements/typecheck.txt
       - .github/workflows/type-checking.yaml

yt_xarray/accessor/_xr_to_yt.py

@@ -494,7 +494,7 @@ def _check_grid_stretchiness(x: npt.NDArray) -> _GridType:
         return _GridType.STRETCHED
 
 
-def _check_for_time(dim_name, dim_vals: np.ndarray):
+def _check_for_time(dim_name, dim_vals: npt.NDArray) -> bool:
     return "time" in dim_name.lower() or type(dim_vals) is np.datetime64
 
 

yt_xarray/accessor/accessor.py

@@ -3,7 +3,7 @@
 import numpy as np
 import xarray as xr
 import yt
-from numpy.typing import ArrayLike
+from numpy.typing import ArrayLike, NDArray
 from unyt import unyt_quantity
 from yt.data_objects.static_output import Dataset as ytDataset
 
@@ -210,7 +210,7 @@ def get_bbox(
         field: str,
         sel_dict: dict[str, Any] | None = None,
         sel_dict_type: str = "isel",
-    ) -> np.ndarray:
+    ) -> NDArray:
         """
         return the bounding box array for a field, with possible selections
 

yt_xarray/utilities/_grid_decomposition.py

@@ -456,28 +456,30 @@ class ChunkInfo:
     def __init__(
         self,
         data_shp: Tuple[int, ...],
-        chunksizes: npt.NDArray,
-        starting_index_offset: npt.NDArray | None = None,
+        chunksizes: npt.NDArray[np.int64],
+        starting_index_offset: npt.NDArray[np.int64] | None = None,
     ):
 
         self.chunksizes = chunksizes
-        self.data_shape = np.asarray(data_shp)
-        self.n_chnk = self.data_shape / chunksizes  # may not be int
+        self.data_shape = np.asarray(data_shp).astype(np.int64)
+        self.n_chnk: npt.NDArray[np.int64] = (
+            self.data_shape / chunksizes
+        )  # may not be int
         self.n_whl_chnk = np.floor(self.n_chnk).astype(int)  # whole chunks in each dim
         self.n_part_chnk = np.ceil(self.n_chnk - self.n_whl_chnk).astype(int)
         self.n_tots = np.prod(self.n_part_chnk + self.n_whl_chnk)
 
         self.ndim = len(data_shp)
         if starting_index_offset is None:
-            starting_index_offset = np.zeros(self.data_shape.shape, dtype=int)
-        self.starting_index_offset = starting_index_offset
+            starting_index_offset = np.zeros(self.data_shape.shape, dtype=np.int64)
+        self.starting_index_offset: npt.NDArray[np.int64] = starting_index_offset
 
-    _si: List[npt.NDArray] | None = None
-    _ei: List[npt.NDArray] | None = None
-    _sizes: List[npt.NDArray] | None = None
+    _si: List[npt.NDArray[np.int64]] | None = None
+    _ei: List[npt.NDArray[np.int64]] | None = None
+    _sizes: List[npt.NDArray[np.int64]] | None = None
 
     @property
-    def si(self) -> List[npt.NDArray]:
+    def si(self) -> List[npt.NDArray[np.int64]]:
         """
         The starting indices of individual chunks by dimension.
         Includes any global offset.
@@ -523,7 +525,7 @@ def si(self) -> List[npt.NDArray]:
         return self._si
 
     @property
-    def ei(self) -> List[npt.NDArray]:
+    def ei(self) -> List[npt.NDArray[np.int64]]:
         """
         The ending indices of individual chunks by dimension.
         Includes any global offset.
@@ -534,7 +536,7 @@ def ei(self) -> List[npt.NDArray]:
         return self._ei
 
     @property
-    def sizes(self) -> List[npt.NDArray]:
+    def sizes(self) -> List[npt.NDArray[np.int64]]:
         if self._sizes is None:
             _ = self.si
             assert self._sizes is not None

yt_xarray/utilities/_utilities.py

@@ -176,7 +176,7 @@ def _test_time_coord(nt: int = 5) -> npt.NDArray:
 
 def _get_test_coord(
     cname, n, minv: float | None = None, maxv: float | None = None
-) -> np.ndarray:
+) -> npt.NDArray[np.floating]:
     if cname in known_coord_aliases:
         cname = known_coord_aliases[cname]
 
@@ -211,7 +211,7 @@ def construct_ds_with_extra_dim(
     ncoords: int | None = None,
     nd_space: int = 3,
     reverse_indices: list[int] | None = None,
-):
+) -> xr.Dataset:
     coord_configs = {
         0: (dim_name, "x", "y", "z"),
         1: (dim_name, "z", "y", "x"),