diff --git a/array_api_strict/_array_object.py b/array_api_strict/_array_object.py index 2d9defa..4d55245 100644 --- a/array_api_strict/_array_object.py +++ b/array_api_strict/_array_object.py @@ -25,6 +25,11 @@ import numpy as np import numpy.typing as npt +try: + from typing import Self +except ImportError: # Python < 3.11 + from typing_extensions import Self + from ._creation_functions import Undef, _undef, asarray from ._devices import CPU_DEVICE, Device, device_supports_dtype from ._dtypes import ( @@ -102,7 +107,7 @@ def _new(cls, x: npt.NDArray[Any] | np.generic, /, device: Device | None) -> Arr return obj # Prevent Array() from working - def __new__(cls, *args: object, **kwargs: object) -> Array: + def __new__(cls, *args: object, **kwargs: object) -> Self: raise TypeError( "The array_api_strict Array object should not be instantiated directly. Use an array creation function, such as asarray(), instead." ) @@ -378,12 +383,10 @@ def _validate_index( _key = key if isinstance(key, tuple) else (key,) for i in _key: if isinstance(i, bool) or not ( - isinstance(i, SupportsIndex) # i.e. ints - or isinstance(i, slice) + isinstance(i, (SupportsIndex, slice)) or i == Ellipsis or (op == "getitem" and i is None) # `None` disallowed in setitem - or isinstance(i, Array) - or isinstance(i, np.ndarray) + or isinstance(i, (Array, np.ndarray)) ): raise IndexError( f"Single-axes index {i} has {type(i)=}, but only " @@ -448,7 +451,7 @@ def _validate_index( else: ellipsis_start = None for pos, i in enumerate(nonexpanding_key): - if not (isinstance(i, Array) or isinstance(i, np.ndarray)): + if not (isinstance(i, (Array, np.ndarray))): if i == Ellipsis: ellipsis_start = pos break @@ -989,7 +992,7 @@ def __xor__(self, other: Array | int, /) -> Array: res = self._array.__xor__(other._array) return self.__class__._new(res, device=self.device) - def __iadd__(self, other: Array | complex, /) -> Array: + def __iadd__(self, other: Array | complex, /) -> Self: """ Performs the operation __iadd__. """ @@ -1012,7 +1015,7 @@ def __radd__(self, other: Array | complex, /) -> Array: res = self._array.__radd__(other._array) return self.__class__._new(res, device=self.device) - def __iand__(self, other: Array | int, /) -> Array: + def __iand__(self, other: Array | int, /) -> Self: """ Performs the operation __iand__. """ @@ -1035,7 +1038,7 @@ def __rand__(self, other: Array | int, /) -> Array: res = self._array.__rand__(other._array) return self.__class__._new(res, device=self.device) - def __ifloordiv__(self, other: Array | float, /) -> Array: + def __ifloordiv__(self, other: Array | float, /) -> Self: """ Performs the operation __ifloordiv__. """ @@ -1058,7 +1061,7 @@ def __rfloordiv__(self, other: Array | float, /) -> Array: res = self._array.__rfloordiv__(other._array) return self.__class__._new(res, device=self.device) - def __ilshift__(self, other: Array | int, /) -> Array: + def __ilshift__(self, other: Array | int, /) -> Self: """ Performs the operation __ilshift__. """ @@ -1081,7 +1084,7 @@ def __rlshift__(self, other: Array | int, /) -> Array: res = self._array.__rlshift__(other._array) return self.__class__._new(res, device=self.device) - def __imatmul__(self, other: Array, /) -> Array: + def __imatmul__(self, other: Array, /) -> Self: """ Performs the operation __imatmul__. """ @@ -1107,7 +1110,7 @@ def __rmatmul__(self, other: Array, /) -> Array: res = self._array.__rmatmul__(other._array) return self.__class__._new(res, device=self.device) - def __imod__(self, other: Array | float, /) -> Array: + def __imod__(self, other: Array | float, /) -> Self: """ Performs the operation __imod__. """ @@ -1130,7 +1133,7 @@ def __rmod__(self, other: Array | float, /) -> Array: res = self._array.__rmod__(other._array) return self.__class__._new(res, device=self.device) - def __imul__(self, other: Array | complex, /) -> Array: + def __imul__(self, other: Array | complex, /) -> Self: """ Performs the operation __imul__. """ @@ -1153,7 +1156,7 @@ def __rmul__(self, other: Array | complex, /) -> Array: res = self._array.__rmul__(other._array) return self.__class__._new(res, device=self.device) - def __ior__(self, other: Array | int, /) -> Array: + def __ior__(self, other: Array | int, /) -> Self: """ Performs the operation __ior__. """ @@ -1176,7 +1179,7 @@ def __ror__(self, other: Array | int, /) -> Array: res = self._array.__ror__(other._array) return self.__class__._new(res, device=self.device) - def __ipow__(self, other: Array | complex, /) -> Array: + def __ipow__(self, other: Array | complex, /) -> Self: """ Performs the operation __ipow__. """ @@ -1201,7 +1204,7 @@ def __rpow__(self, other: Array | complex, /) -> Array: # for 0-d arrays, so we use pow() here instead. return pow(other, self) - def __irshift__(self, other: Array | int, /) -> Array: + def __irshift__(self, other: Array | int, /) -> Self: """ Performs the operation __irshift__. """ @@ -1224,7 +1227,7 @@ def __rrshift__(self, other: Array | int, /) -> Array: res = self._array.__rrshift__(other._array) return self.__class__._new(res, device=self.device) - def __isub__(self, other: Array | complex, /) -> Array: + def __isub__(self, other: Array | complex, /) -> Self: """ Performs the operation __isub__. """ @@ -1247,7 +1250,7 @@ def __rsub__(self, other: Array | complex, /) -> Array: res = self._array.__rsub__(other._array) return self.__class__._new(res, device=self.device) - def __itruediv__(self, other: Array | complex, /) -> Array: + def __itruediv__(self, other: Array | complex, /) -> Self: """ Performs the operation __itruediv__. """ @@ -1270,7 +1273,7 @@ def __rtruediv__(self, other: Array | complex, /) -> Array: res = self._array.__rtruediv__(other._array) return self.__class__._new(res, device=self.device) - def __ixor__(self, other: Array | int, /) -> Array: + def __ixor__(self, other: Array | int, /) -> Self: """ Performs the operation __ixor__. """ diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py index 6d69cb1..f47a136 100644 --- a/array_api_strict/_creation_functions.py +++ b/array_api_strict/_creation_functions.py @@ -36,9 +36,8 @@ def _check_valid_dtype(dtype: DType | None, device: Device | None = None) -> Non if dtype not in _all_dtypes: raise ValueError(f"dtype must be one of the supported dtypes, got {dtype!r}") - if device is not None: - if not device_supports_dtype(device, dtype): - raise ValueError(f"Device {device!r} does not support dtype={dtype!r}.") + if device is not None and not device_supports_dtype(device, dtype): + raise ValueError(f"Device {device!r} does not support dtype={dtype!r}.") def _supports_buffer_protocol(obj: object) -> TypeIs[SupportsBufferProtocol]: @@ -98,9 +97,8 @@ def asarray( if isinstance(obj, Array): return Array._new(np.array(obj._array, copy=copy, dtype=_np_dtype), device=device) - elif isinstance(obj, list | tuple): - if any(isinstance(x, Array) for x in obj): - raise TypeError("Nested Arrays are not allowed. Use `stack` instead.") + elif isinstance(obj, list | tuple) and any(isinstance(x, Array) for x in obj): + raise TypeError("Nested Arrays are not allowed. Use `stack` instead.") if dtype is None and isinstance(obj, int) and (obj > 2 ** 64 or obj < -(2 ** 63)): # Give a better error message in this case. NumPy would convert this diff --git a/array_api_strict/tests/test_flags.py b/array_api_strict/tests/test_flags.py index 52710dd..b834db4 100644 --- a/array_api_strict/tests/test_flags.py +++ b/array_api_strict/tests/test_flags.py @@ -269,10 +269,7 @@ def test_boolean_indexing(): @pytest.mark.parametrize('func_name', linalg_examples.keys()) def test_linalg(func_name): func = linalg_examples[func_name] - if func_name in linalg_main_namespace_examples: - main_namespace_func = linalg_main_namespace_examples[func_name] - else: - main_namespace_func = lambda: None + main_namespace_func = linalg_main_namespace_examples.get(func_name, lambda: None) # First make sure the example actually works func() diff --git a/ruff.toml b/ruff.toml index 15d7ebe..68780a5 100644 --- a/ruff.toml +++ b/ruff.toml @@ -4,7 +4,10 @@ ignore = [ "E402", # Annoying style checks "E7", + # Keep nested with statements "SIM117", + # Keep nested if statements + "SIM102", ] [lint.isort] combine-as-imports = true