diff --git a/coordax/coordinate_systems.py b/coordax/coordinate_systems.py index 9fafc99..764e997 100644 --- a/coordax/coordinate_systems.py +++ b/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,7 +846,7 @@ 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: @@ -854,15 +854,15 @@ def map_indexers_using_ticks( 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,17 +871,17 @@ 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: @@ -889,12 +889,12 @@ def map_indexers_using_ticks( 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) diff --git a/coordax/fields.py b/coordax/fields.py index c9a9325..1b4acdb 100644 --- a/coordax/fields.py +++ b/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: """ 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] ) diff --git a/coordax/named_axes.py b/coordax/named_axes.py index 21c758f..142fcc2 100644 --- a/coordax/named_axes.py +++ b/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__') diff --git a/coordax/ndarrays.py b/coordax/ndarrays.py index 086945a..c4e5af9 100644 --- a/coordax/ndarrays.py +++ b/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 diff --git a/coordax/testing.py b/coordax/testing.py index 98307b7..acab9ca 100644 --- a/coordax/testing.py +++ b/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, diff --git a/docs/conf.py b/docs/conf.py index fe908c1..9e9cb55 100644 --- a/docs/conf.py +++ b/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 ''