ty for type checking

Author: ianhi

pyproject.toml

@@ -33,12 +33,11 @@ dev = [
   "pytest-cov>=7.0.0",
   "ruff>=0.14.8",
   "scipy>=1.16.3",
+  "ty>=0.0.5",
   "xarray-fancy-repr>=0.0.2",
   "zarr>=3.1.5",
 ]
-docs = [
-    "jupyter-book>=2.1.0",
-]
+docs = ["jupyter-book>=2.1.0"]
 vim = ["jupyterlab-vim>=4.1.4"]
 
 [tool.pytest.ini_options]

src/linked_indices/benchmark_utils.py

@@ -15,7 +15,7 @@
 
 
 def timeit_benchmark(
-    stmt: str | callable,
+    stmt: str | Callable[[], Any],
     setup: str = "pass",
     globals: dict[str, Any] | None = None,
     repeat: int = 7,

src/linked_indices/example_data.py

@@ -65,7 +65,7 @@ def speech_annotations() -> pd.DataFrame:
         [4.5, 2.0, "how"],
         [7.0, 2.5, "are"],
     ]
-    return pd.DataFrame(data, columns=["onset", "duration", "word"])
+    return pd.DataFrame(data, columns=["onset", "duration", "word"])  # type: ignore[arg-type]
 
 
 def multi_level_annotations() -> tuple[pd.DataFrame, pd.DataFrame]:
@@ -104,7 +104,8 @@ def multi_level_annotations() -> tuple[pd.DataFrame, pd.DataFrame]:
         [6.0, 3.5, "test", "noun"],
     ]
     word_df = pd.DataFrame(
-        word_data, columns=["onset", "duration", "word", "part_of_speech"]
+        word_data,
+        columns=["onset", "duration", "word", "part_of_speech"],  # type: ignore[arg-type]
     )
 
     # Phoneme-level annotations (more fine-grained)
@@ -118,7 +119,7 @@ def multi_level_annotations() -> tuple[pd.DataFrame, pd.DataFrame]:
         [6.0, 1.2, "t"],
         [7.2, 2.3, "st"],
     ]
-    phoneme_df = pd.DataFrame(phoneme_data, columns=["onset", "duration", "phoneme"])
+    phoneme_df = pd.DataFrame(phoneme_data, columns=["onset", "duration", "phoneme"])  # type: ignore[arg-type]
 
     return word_df, phoneme_df
 
@@ -166,7 +167,7 @@ def mixed_event_annotations() -> pd.DataFrame:
         [0.0, 5.0, "image_A", "stimulus"],
         [5.0, 5.0, "image_B", "stimulus"],
     ]
-    return pd.DataFrame(data, columns=["onset", "duration", "label", "event_type"])
+    return pd.DataFrame(data, columns=["onset", "duration", "label", "event_type"])  # type: ignore[arg-type]
 
 
 def generate_audio_signal(

src/linked_indices/multi_interval_index.py

@@ -1,7 +1,7 @@
 from collections.abc import Mapping
 from dataclasses import dataclass, field
 from numbers import Integral
-from typing import Any
+from typing import Any, cast
 
 from collections import defaultdict
 
@@ -322,6 +322,7 @@ def from_variables(cls, variables, *, options):
                             {name: var}, options=options
                         )
 
+            assert interval_index is not None, f"No interval index found for {dim_name}"
             interval_dims[dim_name] = IntervalDimInfo(
                 dim_name=dim_name,
                 coord_name=coord_name,
@@ -347,8 +348,10 @@ def from_variables(cls, variables, *, options):
             debug=debug,
         )
 
-    def create_variables(self, variables):
-        idx_variables = {}
+    def create_variables(
+        self, variables: Mapping[Any, Variable] | None = None
+    ) -> dict[Any, Variable]:
+        idx_variables: dict[Any, Variable] = {}
 
         idx_variables.update(self._continuous_index.create_variables(variables))
 
@@ -431,7 +434,7 @@ def _get_overlapping_slice(
             )
 
         # Find which intervals overlap
-        overlaps = interval_index.overlaps(query_interval)
+        overlaps = interval_index.overlaps(query_interval)  # type: ignore[union-attr]
         overlap_indices = np.where(overlaps)[0]
 
         if self._debug:
@@ -603,7 +606,7 @@ def isel(
                     continue
 
                 overlap_slice = self._get_overlapping_slice(
-                    info.interval_index.index,
+                    cast(pd.IntervalIndex, info.interval_index.index),
                     time_range,
                 )
 
@@ -681,7 +684,7 @@ def sel(self, labels, method=None, tolerance=None):
             ) is not _MISSING or (
                 dim_name := self._label_to_dim.get(key, _MISSING)
             ) is not _MISSING:
-                info = self._interval_dims[dim_name]
+                info = self._interval_dims[cast(str, dim_name)]
                 # Use label index if key is a label, otherwise use interval index
                 idx = info.label_indexes.get(key, info.interval_index)
                 sel_res = idx.sel({key: value}, method=method, tolerance=tolerance)
@@ -714,7 +717,7 @@ def sel(self, labels, method=None, tolerance=None):
 
                 # Find overlapping intervals
                 overlap_slice = self._get_overlapping_slice(
-                    info.interval_index.index,
+                    cast(pd.IntervalIndex, info.interval_index.index),
                     time_range,
                 )
 

src/linked_indices/nd_index.py

@@ -168,7 +168,7 @@ def from_variables(cls, variables: Mapping[str, Variable], *, options):
 
             nd_coords[name] = NDCoord(
                 name=name,
-                dims=var.dims,
+                dims=tuple(str(d) for d in var.dims),
                 values=var.values,
             )
 
@@ -184,11 +184,13 @@ def from_variables(cls, variables: Mapping[str, Variable], *, options):
 
         return cls(nd_coords=nd_coords, slice_method=slice_method, debug=debug)
 
-    def create_variables(self, variables):
+    def create_variables(
+        self, variables: Mapping[Any, Variable] | None = None
+    ) -> dict[Any, Variable]:
         """Create index variables for the dataset."""
         from xarray.core.variable import Variable as XrVariable
 
-        idx_variables = {}
+        idx_variables: dict[Any, Variable] = {}
         for name, ndc in self._nd_coords.items():
             idx_variables[name] = XrVariable(dims=ndc.dims, data=ndc.values)
         return idx_variables
@@ -236,7 +238,7 @@ def _binary_search(
         self, flat_values: np.ndarray, value: float, method: str | None, coord_name: str
     ) -> int:
         """O(log n) binary search for sorted arrays."""
-        idx = np.searchsorted(flat_values, value)
+        idx = int(np.searchsorted(flat_values, value))
         n = len(flat_values)
 
         if method == "nearest":
@@ -276,7 +278,7 @@ def _find_nearest_index(
         idx = (
             searchsorted_idx
             if searchsorted_idx is not None
-            else np.searchsorted(flat_values, value)
+            else int(np.searchsorted(flat_values, value))
         )
 
         if idx == 0:
@@ -287,9 +289,9 @@ def _find_nearest_index(
             left_val = flat_values[idx - 1]
             right_val = flat_values[idx]
             if abs(value - left_val) <= abs(value - right_val):
-                return idx - 1
+                return int(idx - 1)
             else:
-                return idx
+                return int(idx)
 
     def _linear_search(
         self, values: np.ndarray, value: float, method: str | None, coord_name: str