Adding type suppressions for pyrefly

coordax/coordinate_systems.py

@@ -829,11 +829,11 @@ def map_indexers_using_ticks(
   if axis not in indexers and dim not in indexers:
     return {}, set()
 
-  ticks = ticks if ticks is not None else axis.fields[dim].data
-  if np.unique(ticks).size != ticks.size:
+  ticks = ticks if ticks is not None else axis.fields[dim].data  # pyrefly: ignore[bad-assignment, bad-index]
+  if np.unique(ticks).size != ticks.size:  # pyrefly: ignore[missing-attribute, no-matching-overload]
     raise ValueError(f'Ticks must be unique, got {ticks}')
   if ticks_are_sorted is None:
-    ticks_are_sorted = np.all(ticks[:-1] < ticks[1:])
+    ticks_are_sorted = np.all(ticks[:-1] < ticks[1:])  # pyrefly: ignore[bad-assignment, unsupported-operation]
 
   key = axis if axis in indexers else dim
   assert isinstance(key, (str, Coordinate))  # make pytype happy
@@ -846,23 +846,23 @@ def map_indexers_using_ticks(
           'Indexing with axis requires same type and dims, got index axis'
           f' {value} for slicing {axis=}.'
       )
-    value = value.fields[dim].data
+    value = value.fields[dim].data  # pyrefly: ignore[bad-index]
 
   if isinstance(value, slice):
     if method is not None:
       raise NotImplementedError('Method argument not supported for slices')
 
     start, stop = value.start, value.stop
     if ticks_are_sorted:
-      start_idx, stop_idx = 0, ticks.size
+      start_idx, stop_idx = 0, ticks.size  # pyrefly: ignore[missing-attribute]
       if start is not None:
-        start_idx = np.searchsorted(ticks, start, side='left')
+        start_idx = np.searchsorted(ticks, start, side='left')  # pyrefly: ignore[no-matching-overload]
       if stop is not None:
-        stop_idx = np.searchsorted(ticks, stop, side='right')
+        stop_idx = np.searchsorted(ticks, stop, side='right')  # pyrefly: ignore[no-matching-overload]
 
       return {key: slice(start_idx, stop_idx)}, {key}
     else:
-      mask = np.ones(ticks.size, dtype=bool)
+      mask = np.ones(ticks.size, dtype=bool)  # pyrefly: ignore[missing-attribute]
       if start is not None:
         mask &= ticks >= start
       if stop is not None:
@@ -871,30 +871,30 @@ def map_indexers_using_ticks(
 
   if method == 'nearest':
     if ticks_are_sorted:
-      candidates = np.searchsorted(ticks, value, side='left')
+      candidates = np.searchsorted(ticks, value, side='left')  # pyrefly: ignore[no-matching-overload]
       left = np.maximum(candidates - 1, 0)
-      right = np.minimum(candidates, len(ticks) - 1)
-      d_left = np.abs(value - ticks[left])
-      d_right = np.abs(value - ticks[right])
+      right = np.minimum(candidates, len(ticks) - 1)  # pyrefly: ignore[bad-argument-type]
+      d_left = np.abs(value - ticks[left])  # pyrefly: ignore[unsupported-operation]
+      d_right = np.abs(value - ticks[right])  # pyrefly: ignore[unsupported-operation]
       # In case of ties, prefer the left (smaller) index.
       idx = np.where(d_left <= d_right, left, right)
     elif np.ndim(value) == 0:
       idx = np.abs(ticks - value).argmin()
     else:
-      ticks_view = ticks.reshape((-1,) + (1,) * np.ndim(value))
+      ticks_view = ticks.reshape((-1,) + (1,) * np.ndim(value))  # pyrefly: ignore[missing-attribute]
       idx = np.abs(ticks_view - value).argmin(axis=0)
 
     if np.ndim(idx) == 0:
       idx = int(idx)
     return {key: idx}, {key}
 
   if method is None:
-    sort_indices = None if ticks_are_sorted else np.argsort(ticks)
-    sorted_ticks = ticks if ticks_are_sorted else ticks[sort_indices]
-    indices = np.searchsorted(sorted_ticks, value)
+    sort_indices = None if ticks_are_sorted else np.argsort(ticks)  # pyrefly: ignore[bad-argument-type]
+    sorted_ticks = ticks if ticks_are_sorted else ticks[sort_indices]  # pyrefly: ignore[unsupported-operation]
+    indices = np.searchsorted(sorted_ticks, value)  # pyrefly: ignore[no-matching-overload]
     if sort_indices is not None:
       indices = sort_indices[indices]
-    unique_retrieved = np.sort(np.unique(ticks[indices]))
+    unique_retrieved = np.sort(np.unique(ticks[indices]))  # pyrefly: ignore[unsupported-operation]
     unique_value = np.sort(np.unique(value))
     if unique_retrieved.size != unique_value.size or np.any(
         unique_retrieved != unique_value
@@ -938,7 +938,7 @@ def fields(self) -> dict[str, 'fields.Field']:
 
   @functools.cached_property
   def _sorted_ticks(self) -> bool:
-    return np.all(self.ticks[:-1] <= self.ticks[1:])
+    return np.all(self.ticks[:-1] <= self.ticks[1:])  # pyrefly: ignore[bad-return]
 
   def map_indexers(
       self,
@@ -1201,7 +1201,7 @@ def get_next_match():
     for coord_type in coord_types:
       if coord_type == CartesianProduct or coord_type == Scalar:
         continue
-      result = coord_type.from_xarray(dims, data_array.coords)
+      result = coord_type.from_xarray(dims, data_array.coords)  # pyrefly: ignore[bad-argument-type]
       if isinstance(result, Coordinate):
         return result
       assert isinstance(result, NoCoordinateMatch)

coordax/fields.py

@@ -61,7 +61,7 @@ def _dimension_names(
 ) -> tuple[str | types.EllipsisType, ...]:
   """Returns a tuple of dimension names from a list of names or coordinates."""
   dims_or_name_tuple = lambda x: x.dims if isinstance(x, Coordinate) else (x,)
-  return sum([dims_or_name_tuple(c) for c in names], start=tuple())
+  return sum([dims_or_name_tuple(c) for c in names], start=tuple())  # pyrefly: ignore[bad-return]
 
 
 def _axes_attrs(field: Field) -> str:
@@ -677,7 +677,7 @@ def unwrap(self, *names: str | Coordinate) -> Array:
       ValueError: Field has self.named_dims=('x', 'y') but names=('y', 'x') were
       requested.
     """
-    names = _dimension_names(*names)
+    names = _dimension_names(*names)  # pyrefly: ignore[bad-assignment]
     if names != self.named_dims:
       raise ValueError(
           f'Field has {self.named_dims=} but {names=} were requested.'
@@ -740,7 +740,7 @@ def untag(self, *axis_order: str | Coordinate) -> Field:
     """
     self._validate_matching_coords(axis_order)
     untag_dims = _dimension_names(*axis_order)
-    named_array = self.named_array.untag(*untag_dims)
+    named_array = self.named_array.untag(*untag_dims)  # pyrefly: ignore[bad-argument-type]
     axes = {k: v for k, v in self.axes.items() if k not in untag_dims}
     result = Field.from_namedarray(named_array=named_array, axes=axes)
     return result
@@ -807,7 +807,7 @@ def tag(self, *names: str | Coordinate | ellipsis | None) -> Field:
       :meth:`coordax.Field.untag`
       :func:`coordax.tag`
     """
-    tag_dims = _dimension_names(*names)
+    tag_dims = _dimension_names(*names)  # pyrefly: ignore[bad-argument-type]
     tagged_array = self.named_array.tag(*tag_dims)
     axes = {}
     axes.update(self.axes)
@@ -869,15 +869,16 @@ def broadcast_like(self, other: Self | Coordinate) -> Self:
       <Field dims=('x', 'y') shape=(2, 3) axes={} >
     """
     if isinstance(other, Coordinate):
-      other = shape_struct_field(other)
+      other = shape_struct_field(other)  # pyrefly: ignore[bad-assignment]
     for k, v in self.axes.items():
-      if other.axes.get(k) != v:
+      if other.axes.get(k) != v:  # pyrefly: ignore[missing-attribute]
         raise ValueError(
+            # pyrefly: ignore[missing-attribute]
             'cannot broadcast field because axes corresponding to dimension '
             f'{k!r} do not match: {v} vs {other.axes.get(k)}'
         )
-    return Field.from_namedarray(
-        self.named_array.broadcast_like(other.named_array), other.axes
+    return Field.from_namedarray(  # pyrefly: ignore[bad-return]
+        self.named_array.broadcast_like(other.named_array), other.axes  # pyrefly: ignore[bad-argument-type, missing-attribute]
     )
 
   def isel(
@@ -945,9 +946,9 @@ def isel(
       f = f.tag(tmp_axes[-1])
 
     for dim, indexer in zip(dim_names, indexers.values(), strict=True):
-      post_slice_coord = f.coordinate.isel({dim: indexer})
+      post_slice_coord = f.coordinate.isel({dim: indexer})  # pyrefly: ignore[bad-argument-type]
       data_slice = [slice(None)] * f.ndim
-      data_slice[f.named_axes[dim]] = indexer
+      data_slice[f.named_axes[dim]] = indexer  # pyrefly: ignore[bad-index]
       f = field(f.data[tuple(data_slice)], post_slice_coord)
     return f.untag(*tmp_axes)
 
@@ -1337,7 +1338,7 @@ def contains_dims(
 ) -> bool:
   """Returns True if the field or coordinate contains the given dimensions."""
   c = field_or_coord.coordinate if is_field(field_or_coord) else field_or_coord
-  return coordinate_systems.contains_dims(c, *dims)
+  return coordinate_systems.contains_dims(c, *dims)  # pyrefly: ignore[bad-argument-type]
 
 
 @utils.export
@@ -1382,7 +1383,7 @@ def get_coordinate_part(
   c = field_or_coord.coordinate if is_field(field_or_coord) else field_or_coord
   dim_to_axes = {d: ax for d, ax in zip(c.dims, c.axes)}
   return coordinate_systems.compose(
-      *[d if coordinate_systems.is_coord(d) else dim_to_axes[d] for d in dims]
+      *[d if coordinate_systems.is_coord(d) else dim_to_axes[d] for d in dims]  # pyrefly: ignore[bad-index]
   )
 
 

coordax/named_axes.py

@@ -114,7 +114,7 @@ def get_array_data_with_truncation(
               'Valid mask must be broadcastable to the shape of ``array``, but '
               f'it had extra axis names {bad_names}'
           )
-        mask = mask.order_as(*(d for d in array.dims if d in mask.dims))
+        mask = mask.order_as(*(d for d in array.dims if d in mask.dims))  # pyrefly: ignore[bad-argument-type]
         mask_data = mask.data
       else:
         if np.broadcast_shapes(mask.shape, array.shape) != array.shape:
@@ -126,14 +126,14 @@ def get_array_data_with_truncation(
     if isinstance(array.data, jax.Array):
       return jax_support.JAXArrayAdapter().get_array_data_with_truncation(
           array=array.data,
-          mask=mask_data,
+          mask=mask_data,  # pyrefly: ignore[bad-argument-type]
           edge_items_per_axis=edge_items_per_axis,
       )
     else:
       assert isinstance(array.data, np.ndarray | ndarrays.NDArray)
       return numpy_support.NumpyArrayAdapter().get_array_data_with_truncation(
           array=ndarrays.to_numpy_array(array.data),
-          mask=mask_data,
+          mask=mask_data,  # pyrefly: ignore[bad-argument-type]
           edge_items_per_axis=edge_items_per_axis,
       )
 
@@ -445,7 +445,7 @@ def vectorized_fun(leaf_data):
     def wrap_output(data: Array) -> NamedArray:
       dims = [None] * data.ndim
       for dim, axis in out_axes_dict.items():
-        dims[axis] = dim
+        dims[axis] = dim  # pyrefly: ignore[unsupported-operation]
       return NamedArray(data, tuple(dims))
 
     is_array = lambda x: isinstance(x, Array)
@@ -719,7 +719,7 @@ def tree_unflatten(cls, treedef, leaves: list[Array | object]) -> Self:
 
     # Restored NamedArray objects may have additional or removed leading
     # dimensions, if produced with scan or vmap.
-    result = cls._new_with_padded_or_trimmed_dims(data, dims)
+    result = cls._new_with_padded_or_trimmed_dims(data, dims)  # pyrefly: ignore[bad-argument-type]
     expected_named_shape = _named_shape(dims, shape)
     if result.named_shape != expected_named_shape:
       raise ValueError(
@@ -799,7 +799,7 @@ def tag(self, *dims: str | ellipsis | None) -> Self:
         dim_queue.pop() if dim is None else dim for dim in self.dims
     )
     assert not dim_queue
-    return type(self)(self.data, new_dims)
+    return type(self)(self.data, new_dims)  # pyrefly: ignore[bad-argument-type]
 
   def untag(self, *dims: str) -> Self:
     """Removes the requested dimension names.
@@ -873,10 +873,10 @@ def order_as(self, *dims: str | types.EllipsisType) -> Self:
           dim for dim in self.dims if dim not in explicit_dims
       )
       i = dims.index(...)
-      dims = dims[:i] + implicit_dims + dims[i + 1 :]
+      dims = dims[:i] + implicit_dims + dims[i + 1 :]  # pyrefly: ignore[bad-assignment]
 
     order = tuple(self.dims.index(dim) for dim in dims)
-    return type(self)(self.data.transpose(order), dims)
+    return type(self)(self.data.transpose(order), dims)  # pyrefly: ignore[bad-argument-type]
 
   def broadcast_like(self, other: Self) -> Self:
     """Broadcasts the array to the shape of the other array."""
@@ -902,8 +902,8 @@ def broadcast_like(self, other: Self) -> Self:
   # Convenience wrappers: Elementwise infix operators.
   __lt__ = _nmap_binary_op(operator.lt, 'jax.Array.__lt__')
   __le__ = _nmap_binary_op(operator.le, 'jax.Array.__le__')
-  __eq__ = _nmap_binary_op(operator.eq, 'jax.Array.__eq__')
-  __ne__ = _nmap_binary_op(operator.ne, 'jax.Array.__ne__')
+  __eq__ = _nmap_binary_op(operator.eq, 'jax.Array.__eq__')  # pyrefly: ignore[bad-override]
+  __ne__ = _nmap_binary_op(operator.ne, 'jax.Array.__ne__')  # pyrefly: ignore[bad-override]
   __ge__ = _nmap_binary_op(operator.ge, 'jax.Array.__ge__')
   __gt__ = _nmap_binary_op(operator.gt, 'jax.Array.__gt__')
 

coordax/ndarrays.py

@@ -141,8 +141,8 @@ def register_ndarray(
     :class:`coordax.experimental.NDArray`
   """
   NDArray.register(array_type)
-  _TO_NUMPY_FUNCS.append((array_type, to_numpy))
-  _FROM_NUMPY_FUNCS.append((is_matching_numpy_array, from_numpy))
+  _TO_NUMPY_FUNCS.append((array_type, to_numpy))  # pyrefly: ignore[bad-argument-type]
+  _FROM_NUMPY_FUNCS.append((is_matching_numpy_array, from_numpy))  # pyrefly: ignore[bad-argument-type]
   return array_type
 
 

coordax/testing.py

@@ -52,7 +52,7 @@ def assert_fields_allclose(
   assert_field_properties(
       actual=actual,
       data=desired.data,
-      dims=desired.dims,
+      dims=desired.dims,  # pyrefly: ignore[bad-argument-type]
       shape=desired.shape,
       axes=desired.axes,
       named_shape=desired.named_shape,
@@ -67,7 +67,7 @@ def assert_fields_equal(actual: coordax.Field, desired: coordax.Field):
   assert_field_properties(
       actual=actual,
       data=desired.data,
-      dims=desired.dims,
+      dims=desired.dims,  # pyrefly: ignore[bad-argument-type]
       shape=desired.shape,
       axes=desired.axes,
       named_shape=desired.named_shape,

docs/conf.py

@@ -128,7 +128,7 @@ def linkcode_resolve(domain, info):
   except Exception as e:
     print(f'did not find source code for: {info}: {e}')
     return None
-  filename = os.path.relpath(
+  filename = os.path.relpath(  # pyrefly: ignore[no-matching-overload]
       filename, start=os.path.dirname(coordax.__file__)
   )
   lines = f'#L{linenum}-L{linenum + len(source)}' if linenum else ''