Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 22 additions & 19 deletions array_api_strict/_array_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment thread
ev-br marked this conversation as resolved.

from ._creation_functions import Undef, _undef, asarray
from ._devices import CPU_DEVICE, Device, device_supports_dtype
from ._dtypes import (
Expand Down Expand Up @@ -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."
)
Expand Down Expand Up @@ -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 "
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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__.
"""
Expand All @@ -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__.
"""
Expand All @@ -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__.
"""
Expand All @@ -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__.
"""
Expand All @@ -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__.
"""
Expand All @@ -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__.
"""
Expand All @@ -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__.
"""
Expand All @@ -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__.
"""
Expand All @@ -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__.
"""
Expand All @@ -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__.
"""
Expand All @@ -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__.
"""
Expand All @@ -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__.
"""
Expand All @@ -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__.
"""
Expand Down
10 changes: 4 additions & 6 deletions array_api_strict/_creation_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down Expand Up @@ -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
Expand Down
5 changes: 1 addition & 4 deletions array_api_strict/tests/test_flags.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
3 changes: 3 additions & 0 deletions ruff.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading