another batch

Author: chrishavlin

.github/workflows/type-check.yml

@@ -8,6 +8,7 @@ on:
     paths:
       - yt_experiments/**/*.py
       - pyproject.toml
+      - requirements/typecheck.txt
       - .github/workflows/type-checking.yaml
   workflow_dispatch:
 
@@ -17,11 +18,6 @@ jobs:
     name: type check
     timeout-minutes: 60
 
-    concurrency:
-      # auto-cancel any in-progress job *on the same branch*
-      group: ${{ github.workflow }}-${{ github.ref }}
-      cancel-in-progress: true
-
     steps:
     - name: Checkout repo
       uses: actions/checkout@v4
@@ -37,7 +33,7 @@ jobs:
     - name: Build
       run: |
         python3 -m pip install --upgrade pip
-        python3 -m pip install "mypy==1.11.2"
+        python3 -m pip install -r requirements/typecheck.txt
 
     - name: list installed deps
       run: python -m pip list

pyproject.toml

@@ -54,3 +54,4 @@ warn_unreachable = true
 disallow_untyped_defs = false
 disallow_incomplete_defs = false
 disable_error_code = ["import-untyped", "import-not-found"]
+no_implicit_reexport = false

requirements/typecheck.txt

@@ -0,0 +1 @@
+mypy==1.11.2

yt_xarray/accessor/_xr_to_yt.py

@@ -45,14 +45,14 @@ def __init__(
         self.sel_dict_type = sel_dict_type
 
         # all these attributes are set in _process_selection
-        self.selected_shape: Tuple[int] | None = None
-        self.full_bbox = None
+        self.selected_shape: Tuple[int, ...]
+        self.full_bbox: npt.NDArray
         self.selected_bbox: npt.NDArray
-        self.full_coords: Tuple[str]
-        self.selected_coords: Tuple[str]
+        self.full_coords: Tuple[str, ...]
+        self.selected_coords: Tuple[str, ...]
         self.starting_indices: npt.NDArray
-        self.selected_time: float | None = None
-        self.ndims: int | None = None
+        self.selected_time: float
+        self.ndims: int
         self.grid_type: _GridType
         self.cell_widths: list[Any]
         self.global_dims: list[Any]
@@ -274,7 +274,9 @@ def select_from_xr(self, xr_ds: xr.Dataset, field: str) -> xr.DataArray:
 
         return vars
 
-    def interp_validation(self, geometry):
+    def interp_validation(
+        self, geometry: str
+    ) -> tuple[bool, tuple[int, ...], npt.NDArray]:
         # checks if yt will need to interpolate to cell center
         # returns a tuple of (bool, shape, bbox). If the bool is True then
         # interpolation is required.
@@ -345,7 +347,7 @@ def interp_validation(self, geometry):
 known_coord_aliases = _default_known_coord_aliases.copy()
 
 
-def reset_coordinate_aliases():
+def reset_coordinate_aliases() -> None:
     kys_to_pop = [
         ky
         for ky in known_coord_aliases.keys()
@@ -399,7 +401,9 @@ def _cf_xr_coord_disamb(
     return None, True
 
 
-def _convert_to_yt_internal_coords(coord_list: tuple[str], xr_field: xr.DataArray):
+def _convert_to_yt_internal_coords(
+    coord_list: tuple[str] | list[str], xr_field: xr.DataArray
+):
     yt_coords = []
     for c in coord_list:
         cname = c.lower()

yt_xarray/sample_data.py

@@ -1,10 +1,12 @@
+from typing import Mapping
+
 import numpy as np
 import xarray as xr
 
 
 def load_random_xr_data(
-    fields: dict[str, tuple[str, ...]],
-    dims: dict[str, tuple[int | float, int | float, int]],
+    fields: Mapping[str, tuple[str, ...]],
+    dims: Mapping[str, tuple[int | float, int | float, int]],
     length_unit: str | None = None,
 ) -> xr.Dataset:
     """

yt_xarray/tests/test_chunking.py

@@ -142,7 +142,7 @@ def test_chunk_bad_length():
         "field0": ("x", "y", "z"),
         "field1": ("x", "y", "z"),
     }
-    dims = {"x": (0, 1, 30), "y": (0, 2, 40), "z": (-1, 0.5, 20)}
+    dims = {"x": (0, 1, 30), "y": (0, 2, 40), "z": (-1.0, 0.5, 20)}
     ds = sample_data.load_random_xr_data(fields, dims)
 
     with pytest.raises(ValueError, match="The number of elements in "):
@@ -175,7 +175,7 @@ def test_chunk_info_caching():
     chunksizes = np.array([5, 5, 5], dtype="int")
     data_shape = (10, 15, 20)
 
-    def _get_ch():
+    def _get_ch() -> ChunkInfo:
         return ChunkInfo(data_shape, chunksizes)
 
     ch = _get_ch()

yt_xarray/tests/test_xr_to_yt.py

@@ -1,9 +1,12 @@
 import builtins
+from typing import Any
 
 import numpy as np
 import pytest
 import xarray as xr
 import yt
+from numpy import typing as npt
+from yt.data_objects.static_output import Dataset as ytDataset
 
 import yt_xarray.accessor._xr_to_yt as xr2yt
 from yt_xarray.utilities._utilities import (
@@ -39,7 +42,7 @@ def ds_xr():
 def test_selection_aliases(coord):
     for othername in xr2yt._coord_aliases[coord]:
         kwargs = {c_m_ds_kwargs[coord]: othername}
-        ds = construct_minimal_ds(**kwargs)
+        ds = construct_minimal_ds(**kwargs)  # type: ignore[arg-type]
         fields = list(ds.data_vars)
         sel = xr2yt.Selection(ds, fields)
         assert np.all(sel.starting_indices == np.array((0, 0, 0)))
@@ -54,7 +57,13 @@ def test_selection_aliases(coord):
             assert xr2yt.known_coord_aliases[othername] in sel.yt_coord_names
 
 
-def _isel_tester(ds_xr, sel, fields, coord, start_index):
+def _isel_tester(
+    ds_xr: xr.Dataset,
+    sel: xr2yt.Selection,
+    fields: list[str],
+    coord: str,
+    start_index: int,
+) -> None:
     dim_id = ds_xr.data_vars[fields[0]].dims.index(coord)
     expected = np.array((0, 0, 0))
     expected[dim_id] = start_index
@@ -73,9 +82,11 @@ def _isel_tester(ds_xr, sel, fields, coord, start_index):
 
 
 @pytest.mark.parametrize("coord", ("latitude", "longitude", "depth"))
-def test_selection_isel(ds_xr, coord):
+def test_selection_isel(ds_xr, coord_input):
+    coord: str = coord_input
     fields = list(ds_xr.data_vars)
 
+    sel_dict: dict[str, Any]
     sel_dict = {coord: slice(1, len(ds_xr.coords[coord]))}
     sel_dict_type = "isel"
     sel = xr2yt.Selection(ds_xr, fields, sel_dict=sel_dict, sel_dict_type=sel_dict_type)
@@ -476,7 +487,9 @@ def test_add_3rd_axis_name(yt_geom):
         _ = xr2yt._add_3rd_axis_name("bad_geometry", expected[:-1])
 
 
-def _get_pixelized_slice(yt_ds):
+def _get_pixelized_slice(
+    yt_ds: ytDataset,
+) -> tuple[Any, tuple[npt.NDArray, npt.NDArray | None] | npt.NDArray]:
     slc = yt_ds.slice(
         yt_ds.coordinates.axis_id["depth"],
         yt_ds.domain_center[yt_ds.coordinates.axis_id["depth"]],
@@ -492,7 +505,9 @@ def _get_pixelized_slice(yt_ds):
     return slc, vals
 
 
-def _get_ds_for_reverse_tests(stretched, use_callable, chunksizes):
+def _get_ds_for_reverse_tests(
+    stretched: bool, use_callable: bool, chunksizes: tuple[int, ...]
+) -> ytDataset:
     ds = construct_minimal_ds(
         min_x=1,
         max_x=359,

yt_xarray/transformations.py

@@ -1,10 +1,11 @@
 import abc
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable, List, Mapping, Optional, Tuple, Union
 
 import numpy as np
 import unyt
 import xarray as xr
 import yt
+from numpy import typing as npt
 from unyt import earth_radius as _earth_radius
 
 from yt_xarray.accessor import _xr_to_yt
@@ -119,7 +120,7 @@ def _validate_input_coords(self, coords, input_coord_type: str):
 
         return new_coords
 
-    def to_native(self, **coords):
+    def to_native(self, **coords: npt.NDArray) -> list[npt.NDArray]:
         """
         Calculate the native coordinates from transformed coordinates.
 
@@ -131,7 +132,7 @@ def to_native(self, **coords):
 
         Returns
         -------
-        tuple
+        list
             coordinate values in the native coordinate system, in order
             of the native_coords attribute.
 
@@ -144,7 +145,7 @@ def to_native(self, **coords):
         new_coords = self._validate_input_coords(coords, "transformed")
         return self._calculate_native(**new_coords)
 
-    def to_transformed(self, **coords):
+    def to_transformed(self, **coords: npt.NDArray) -> list[npt.NDArray]:
         """
         Calculate the transformed coordinates from native coordinates.
 
@@ -156,7 +157,7 @@ def to_transformed(self, **coords):
 
         Returns
         -------
-        tuple
+        list
             coordinate values in the transformed coordinate system, in order
             of the transformed_coords attribute.
 
@@ -223,7 +224,7 @@ class LinearScale(Transformer):
 
     """
 
-    def __init__(self, native_coords: Tuple[str], scale: Optional[dict] = None):
+    def __init__(self, native_coords: Tuple[str, ...], scale: Optional[dict] = None):
         if scale is None:
             scale = {}
 
@@ -234,20 +235,22 @@ def __init__(self, native_coords: Tuple[str], scale: Optional[dict] = None):
         transformed_coords = tuple([nc + "_sc" for nc in native_coords])
         super().__init__(native_coords, transformed_coords)
 
-    def _calculate_transformed(self, **coords):
+    def _calculate_transformed(self, **coords) -> list[npt.NDarray]:
         transformed = []
         for nc_sc in self.transformed_coords:
             nc = nc_sc[:-3]  # native coord name. e.g., go from "x_sc" to just "x"
             transformed.append(np.asarray(coords[nc]) * self.scale[nc])
         return transformed
 
-    def _calculate_native(self, **coords):
+    def _calculate_native(self, **coords) -> list[npt.NDarray]:
         native = []
         for nc in self.native_coords:
             native.append(np.asarray(coords[nc + "_sc"]) / self.scale[nc])
         return native
 
-    def calculate_transformed_bbox(self, bbox_dict: dict) -> np.ndarray:
+    def calculate_transformed_bbox(
+        self, bbox_dict: Mapping[str, npt.NDArray]
+    ) -> npt.NDArray:
         """
         Calculates a bounding box in transformed coordinates for a bounding box dictionary
         in native coordinates.

yt_xarray/utilities/_grid_decomposition.py

@@ -442,7 +442,7 @@ class ChunkInfo:
 
     def __init__(
         self,
-        data_shp: Tuple[int,],
+        data_shp: Tuple[int, ...],
         chunksizes: np.ndarray,
         starting_index_offset: np.ndarray = None,
     ):

yt_xarray/utilities/_utilities.py

@@ -242,7 +242,7 @@ def construct_ds_with_extra_dim(
     return xr.Dataset(data_vars=data_vars)
 
 
-def _find_file(file: PathLike[str]) -> PathLike[str] | str:
+def _find_file(file: PathLike[str] | str) -> PathLike[str] | str:
     if os.path.isfile(file):
         return file