array-api-strict 2.3__py3-none-any.whl → 2.4__py3-none-any.whl

array_api_strict/_constants.py

@@ -4,4 +4,4 @@ e = np.e
 inf = np.inf
 nan = np.nan
 pi = np.pi
-newaxis = np.newaxis
+newaxis: None = np.newaxis

array_api_strict/_creation_functions.py

@@ -1,51 +1,47 @@
 from __future__ import annotations
 
-from contextlib import contextmanager
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+from enum import Enum
+from typing import TYPE_CHECKING, Literal
 
-if TYPE_CHECKING:
-    from ._typing import (
-        Array,
-        Device,
-        Dtype,
-        NestedSequence,
-        SupportsBufferProtocol,
-    )
-from ._dtypes import _DType, _all_dtypes
+import numpy as np
+
+from ._dtypes import DType, _all_dtypes, _np_dtype
 from ._flags import get_array_api_strict_flags
+from ._typing import NestedSequence, SupportsBufferProtocol, SupportsDLPack
 
-import numpy as np
+if TYPE_CHECKING:
+    # TODO import from typing (requires Python >=3.13)
+    from typing_extensions import TypeIs
 
-@contextmanager
-def allow_array():
-    """
-    Temporarily enable Array.__array__. This is needed for np.array to parse
-    list of lists of Array objects.
-    """
-    from . import _array_object
-    original_value = _array_object._allow_array
-    try:
-        _array_object._allow_array = True
-        yield
-    finally:
-        _array_object._allow_array = original_value
+    # Circular import
+    from ._array_object import Array, Device
 
-def _check_valid_dtype(dtype):
+
+class Undef(Enum):
+    UNDEF = 0
+
+
+_undef = Undef.UNDEF
+
+
+def _check_valid_dtype(dtype: DType | None) -> None:
     # Note: Only spelling dtypes as the dtype objects is supported.
     if dtype not in (None,) + _all_dtypes:
         raise ValueError(f"dtype must be one of the supported dtypes, got {dtype!r}")
 
-def _supports_buffer_protocol(obj):
+
+def _supports_buffer_protocol(obj: object) -> TypeIs[SupportsBufferProtocol]:
     try:
-        memoryview(obj)
+        memoryview(obj)  # type: ignore[arg-type]
     except TypeError:
         return False
     return True
 
-def _check_device(device):
+
+def _check_device(device: Device | None) -> None:
     # _array_object imports in this file are inside the functions to avoid
     # circular imports
-    from ._array_object import Device, ALL_DEVICES
+    from ._array_object import ALL_DEVICES, Device
 
     if device is not None and not isinstance(device, Device):
         raise ValueError(f"Unsupported device {device!r}")
@@ -53,20 +49,14 @@ def _check_device(device):
     if device is not None and device not in ALL_DEVICES:
         raise ValueError(f"Unsupported device {device!r}")
 
+
 def asarray(
-    obj: Union[
-        Array,
-        bool,
-        int,
-        float,
-        NestedSequence[bool | int | float],
-        SupportsBufferProtocol,
-    ],
+    obj: Array | complex | NestedSequence[complex] | SupportsBufferProtocol,
     /,
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
-    copy: Optional[bool] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
+    copy: bool | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.asarray <numpy.asarray>`.
@@ -108,23 +98,27 @@ 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.")
+
     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
         # to an object array. TODO: This won't handle large integers in lists.
         raise OverflowError("Integer out of bounds for array dtypes")
-    with allow_array():
-        res = np.array(obj, dtype=_np_dtype, copy=copy)
+
+    res = np.array(obj, dtype=_np_dtype, copy=copy)
     return Array._new(res, device=device)
 
 
 def arange(
-    start: Union[int, float],
+    start: float,
     /,
-    stop: Optional[Union[int, float]] = None,
-    step: Union[int, float] = 1,
+    stop: float | None = None,
+    step: float = 1,
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.arange <numpy.arange>`.
@@ -136,16 +130,17 @@ def arange(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.arange(start, stop=stop, step=step, dtype=dtype), device=device)
+    return Array._new(
+        np.arange(start, stop, step, dtype=_np_dtype(dtype)),
+        device=device,
+    )
 
 
 def empty(
-    shape: Union[int, Tuple[int, ...]],
+    shape: int | tuple[int, ...],
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.empty <numpy.empty>`.
@@ -157,13 +152,11 @@ def empty(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.empty(shape, dtype=dtype), device=device)
+    return Array._new(np.empty(shape, dtype=_np_dtype(dtype)), device=device)
 
 
 def empty_like(
-    x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
+    x: Array, /, *, dtype: DType | None = None, device: Device | None = None
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.empty_like <numpy.empty_like>`.
@@ -177,19 +170,17 @@ def empty_like(
     if device is None:
         device = x.device
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.empty_like(x._array, dtype=dtype), device=device)
+    return Array._new(np.empty_like(x._array, dtype=_np_dtype(dtype)), device=device)
 
 
 def eye(
     n_rows: int,
-    n_cols: Optional[int] = None,
+    n_cols: int | None = None,
     /,
     *,
     k: int = 0,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.eye <numpy.eye>`.
@@ -201,45 +192,43 @@ def eye(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.eye(n_rows, M=n_cols, k=k, dtype=dtype), device=device)
-
+    return Array._new(
+        np.eye(n_rows, M=n_cols, k=k, dtype=_np_dtype(dtype)), device=device
+    )
 
-_default = object()
 
 def from_dlpack(
-    x: object,
+    x: SupportsDLPack,
     /,
     *,
-    device: Optional[Device] = _default,
-    copy: Optional[bool] = _default,
+    device: Device | Undef | None = _undef,
+    copy: bool | Undef | None = _undef,
 ) -> Array:
     from ._array_object import Array
 
     if get_array_api_strict_flags()['api_version'] < '2023.12':
-        if device is not _default:
+        if device is not _undef:
             raise ValueError("The device argument to from_dlpack requires at least version 2023.12 of the array API")
-        if copy is not _default:
+        if copy is not _undef:
             raise ValueError("The copy argument to from_dlpack requires at least version 2023.12 of the array API")
 
     # Going to wait for upstream numpy support
-    if device is not _default:
+    if device is not _undef:
         _check_device(device)
     else:
         device = None
-    if copy not in [_default, None]:
+    if copy not in [_undef, None]:
         raise NotImplementedError("The copy argument to from_dlpack is not yet implemented")
 
     return Array._new(np.from_dlpack(x), device=device)
 
 
 def full(
-    shape: Union[int, Tuple[int, ...]],
-    fill_value: Union[int, float],
+    shape: int | tuple[int, ...],
+    fill_value: complex,
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.full <numpy.full>`.
@@ -253,10 +242,8 @@ def full(
 
     if isinstance(fill_value, Array) and fill_value.ndim == 0:
         fill_value = fill_value._array
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    res = np.full(shape, fill_value, dtype=dtype)
-    if _DType(res.dtype) not in _all_dtypes:
+    res = np.full(shape, fill_value, dtype=_np_dtype(dtype))
+    if DType(res.dtype) not in _all_dtypes:
         # This will happen if the fill value is not something that NumPy
         # coerces to one of the acceptable dtypes.
         raise TypeError("Invalid input to full")
@@ -266,10 +253,10 @@ def full(
 def full_like(
     x: Array,
     /,
-    fill_value: Union[int, float],
+    fill_value: complex,
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.full_like <numpy.full_like>`.
@@ -283,10 +270,8 @@ def full_like(
     if device is None:
         device = x.device
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    res = np.full_like(x._array, fill_value, dtype=dtype)
-    if _DType(res.dtype) not in _all_dtypes:
+    res = np.full_like(x._array, fill_value, dtype=_np_dtype(dtype))
+    if DType(res.dtype) not in _all_dtypes:
         # This will happen if the fill value is not something that NumPy
         # coerces to one of the acceptable dtypes.
         raise TypeError("Invalid input to full_like")
@@ -294,13 +279,13 @@ def full_like(
 
 
 def linspace(
-    start: Union[int, float],
-    stop: Union[int, float],
+    start: complex,
+    stop: complex,
     /,
     num: int,
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
     endpoint: bool = True,
 ) -> Array:
     """
@@ -313,12 +298,13 @@ def linspace(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint), device=device)
+    return Array._new(
+        np.linspace(start, stop, num, dtype=_np_dtype(dtype), endpoint=endpoint),
+        device=device,
+    )
 
 
-def meshgrid(*arrays: Array, indexing: str = "xy") -> List[Array]:
+def meshgrid(*arrays: Array, indexing: Literal["xy", "ij"] = "xy") -> list[Array]:
     """
     Array API compatible wrapper for :py:func:`np.meshgrid <numpy.meshgrid>`.
 
@@ -348,10 +334,10 @@ def meshgrid(*arrays: Array, indexing: str = "xy") -> List[Array]:
 
 
 def ones(
-    shape: Union[int, Tuple[int, ...]],
+    shape: int | tuple[int, ...],
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.ones <numpy.ones>`.
@@ -363,13 +349,11 @@ def ones(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.ones(shape, dtype=dtype), device=device)
+    return Array._new(np.ones(shape, dtype=_np_dtype(dtype)), device=device)
 
 
 def ones_like(
-    x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
+    x: Array, /, *, dtype: DType | None = None, device: Device | None = None
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.ones_like <numpy.ones_like>`.
@@ -383,9 +367,7 @@ def ones_like(
     if device is None:
         device = x.device
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.ones_like(x._array, dtype=dtype), device=device)
+    return Array._new(np.ones_like(x._array, dtype=_np_dtype(dtype)), device=device)
 
 
 def tril(x: Array, /, *, k: int = 0) -> Array:
@@ -417,10 +399,10 @@ def triu(x: Array, /, *, k: int = 0) -> Array:
 
 
 def zeros(
-    shape: Union[int, Tuple[int, ...]],
+    shape: int | tuple[int, ...],
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
+    dtype: DType | None = None,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.zeros <numpy.zeros>`.
@@ -432,13 +414,11 @@ def zeros(
     _check_valid_dtype(dtype)
     _check_device(device)
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.zeros(shape, dtype=dtype), device=device)
+    return Array._new(np.zeros(shape, dtype=_np_dtype(dtype)), device=device)
 
 
 def zeros_like(
-    x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
+    x: Array, /, *, dtype: DType | None = None, device: Device | None = None
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.zeros_like <numpy.zeros_like>`.
@@ -452,6 +432,4 @@ def zeros_like(
     if device is None:
         device = x.device
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-    return Array._new(np.zeros_like(x._array, dtype=dtype), device=device)
+    return Array._new(np.zeros_like(x._array, dtype=_np_dtype(dtype)), device=device)

array_api_strict/_data_type_functions.py

@@ -1,38 +1,35 @@
-from __future__ import annotations
+from dataclasses import dataclass
+
+import numpy as np
 
-from ._array_object import Array
-from ._creation_functions import _check_device
+from ._array_object import Array, Device
+from ._creation_functions import Undef, _check_device, _undef
 from ._dtypes import (
-    _DType,
+    DType,
     _all_dtypes,
     _boolean_dtypes,
-    _signed_integer_dtypes,
-    _unsigned_integer_dtypes,
-    _integer_dtypes,
-    _real_floating_dtypes,
     _complex_floating_dtypes,
+    _integer_dtypes,
     _numeric_dtypes,
+    _real_floating_dtypes,
     _result_type,
+    _signed_integer_dtypes,
+    _unsigned_integer_dtypes,
 )
 from ._flags import get_array_api_strict_flags
 
-from dataclasses import dataclass
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from typing import List, Tuple, Union, Optional
-    from ._typing import Dtype, Device
-
-import numpy as np
-
-# Use to emulate the asarray(device) argument not existing in 2022.12
-_default = object()
 
 # Note: astype is a function, not an array method as in NumPy.
 def astype(
-    x: Array, dtype: Dtype, /, *, copy: bool = True, device: Optional[Device] = _default
+    x: Array,
+    dtype: DType,
+    /,
+    *,
+    copy: bool = True,
+    # _default is used to emulate the device argument not existing in 2022.12
+    device: Device | Undef | None = _undef,
 ) -> Array:
-    if device is not _default:
+    if device is not _undef:
         if get_array_api_strict_flags()['api_version'] >= '2023.12':
             _check_device(device)
         else:
@@ -52,7 +49,7 @@ def astype(
     return Array._new(x._array.astype(dtype=dtype._np_dtype, copy=copy), device=device)
 
 
-def broadcast_arrays(*arrays: Array) -> List[Array]:
+def broadcast_arrays(*arrays: Array) -> list[Array]:
     """
     Array API compatible wrapper for :py:func:`np.broadcast_arrays <numpy.broadcast_arrays>`.
 
@@ -65,7 +62,7 @@ def broadcast_arrays(*arrays: Array) -> List[Array]:
     ]
 
 
-def broadcast_to(x: Array, /, shape: Tuple[int, ...]) -> Array:
+def broadcast_to(x: Array, /, shape: tuple[int, ...]) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.broadcast_to <numpy.broadcast_to>`.
 
@@ -76,7 +73,7 @@ def broadcast_to(x: Array, /, shape: Tuple[int, ...]) -> Array:
     return Array._new(np.broadcast_to(x._array, shape), device=x.device)
 
 
-def can_cast(from_: Union[Dtype, Array], to: Dtype, /) -> bool:
+def can_cast(from_: DType | Array, to: DType, /) -> bool:
     """
     Array API compatible wrapper for :py:func:`np.can_cast <numpy.can_cast>`.
 
@@ -103,7 +100,9 @@ def can_cast(from_: Union[Dtype, Array], to: Dtype, /) -> bool:
 
 # These are internal objects for the return types of finfo and iinfo, since
 # the NumPy versions contain extra data that isn't part of the spec.
-@dataclass
+# There should be no expectation for them to be comparable, hashable, or writeable.
+
+@dataclass(frozen=True, eq=False)
 class finfo_object:
     bits: int
     # Note: The types of the float data here are float, whereas in NumPy they
@@ -112,26 +111,35 @@ class finfo_object:
     max: float
     min: float
     smallest_normal: float
-    dtype: Dtype
+    dtype: DType
+
+    __hash__ = NotImplemented
 
 
-@dataclass
+@dataclass(frozen=True, eq=False)
 class iinfo_object:
     bits: int
     max: int
     min: int
-    dtype: Dtype
+    dtype: DType
 
+    __hash__ = NotImplemented
 
-def finfo(type: Union[Dtype, Array], /) -> finfo_object:
+
+def finfo(type: DType | Array, /) -> finfo_object:
     """
     Array API compatible wrapper for :py:func:`np.finfo <numpy.finfo>`.
 
     See its docstring for more information.
     """
-    if isinstance(type, _DType):
-        type = type._np_dtype
-    fi = np.finfo(type)
+    if isinstance(type, Array):
+        np_type = type._dtype._np_dtype
+    elif isinstance(type, DType):
+        np_type = type._np_dtype
+    else:
+        raise TypeError(f"'type' must be a dtype or array, not {type!r}")
+
+    fi = np.finfo(np_type)
     # Note: The types of the float data here are float, whereas in NumPy they
     # are scalars of the corresponding float dtype.
     return finfo_object(
@@ -140,35 +148,39 @@ def finfo(type: Union[Dtype, Array], /) -> finfo_object:
         float(fi.max),
         float(fi.min),
         float(fi.smallest_normal),
-        fi.dtype,
+        DType(fi.dtype),
     )
 
 
-def iinfo(type: Union[Dtype, Array], /) -> iinfo_object:
+def iinfo(type: DType | Array, /) -> iinfo_object:
     """
     Array API compatible wrapper for :py:func:`np.iinfo <numpy.iinfo>`.
 
     See its docstring for more information.
     """
-    if isinstance(type, _DType):
-        type = type._np_dtype
-    ii = np.iinfo(type)
-    return iinfo_object(ii.bits, ii.max, ii.min, ii.dtype)
+    if isinstance(type, Array):
+        np_type = type._dtype._np_dtype
+    elif isinstance(type, DType):
+        np_type = type._np_dtype
+    else:
+        raise TypeError(f"'type' must be a dtype or array, not {type!r}")
+
+    ii = np.iinfo(np_type)
+    return iinfo_object(ii.bits, ii.max, ii.min, DType(ii.dtype))
 
 
 # Note: isdtype is a new function from the 2022.12 array API specification.
-def isdtype(
-    dtype: Dtype, kind: Union[Dtype, str, Tuple[Union[Dtype, str], ...]]
-) -> bool:
+def isdtype(dtype: DType, kind: DType | str | tuple[DType | str, ...]) -> bool:
     """
-    Returns a boolean indicating whether a provided dtype is of a specified data type ``kind``.
+    Returns a boolean indicating whether a provided dtype is of a specified
+    data type ``kind``.
 
     See
     https://data-apis.org/array-api/latest/API_specification/generated/array_api.isdtype.html
     for more details
     """
-    if not isinstance(dtype, _DType):
-       raise TypeError(f"'dtype' must be a dtype, not a {type(dtype)!r}")
+    if not isinstance(dtype, DType):
+        raise TypeError(f"'dtype' must be a dtype, not a {type(dtype)!r}")
 
     if isinstance(kind, tuple):
         # Disallow nested tuples
@@ -197,7 +209,8 @@ def isdtype(
     else:
         raise TypeError(f"'kind' must be a dtype, str, or tuple of dtypes and strs, not {type(kind).__name__}")
 
-def result_type(*arrays_and_dtypes: Union[Array, Dtype, int, float, complex, bool]) -> Dtype:
+
+def result_type(*arrays_and_dtypes: Array | DType | complex) -> DType:
     """
     Array API compatible wrapper for :py:func:`np.result_type <numpy.result_type>`.
 
@@ -207,35 +220,31 @@ def result_type(*arrays_and_dtypes: Union[Array, Dtype, int, float, complex, boo
     # required by the spec rather than using np.result_type. NumPy implements
     # too many extra type promotions like int64 + uint64 -> float64, and does
     # value-based casting on scalar arrays.
-    A = []
+    dtypes = []
     scalars = []
     for a in arrays_and_dtypes:
-        if isinstance(a, Array):
-            a = a.dtype
+        if isinstance(a, DType):
+            dtypes.append(a)
+        elif isinstance(a, Array):
+            dtypes.append(a.dtype)
         elif isinstance(a, (bool, int, float, complex)):
             scalars.append(a)
-        elif isinstance(a, np.ndarray) or a not in _all_dtypes:
-            raise TypeError("result_type() inputs must be array_api arrays or dtypes")
-        A.append(a)
-
-    # remove python scalars
-    A = [a for a in A if not isinstance(a, (bool, int, float, complex))]
+        else:
+            raise TypeError(
+                "result_type() inputs must be Array API arrays, dtypes, or scalars"
+            )
 
-    if len(A) == 0:
+    if not dtypes:
         raise ValueError("at least one array or dtype is required")
-    elif len(A) == 1:
-        result = A[0]
-    else:
-        t = A[0]
-        for t2 in A[1:]:
-            t = _result_type(t, t2)
-        result = t
+    result = dtypes[0]
+    for t2 in dtypes[1:]:
+        result = _result_type(result, t2)
 
-    if len(scalars) == 0:
+    if not scalars:
         return result
 
     if get_array_api_strict_flags()['api_version'] <= '2023.12':
-        raise TypeError("result_type() inputs must be array_api arrays or dtypes")
+        raise TypeError("result_type() inputs must be Array API arrays or dtypes")
 
     # promote python scalars given the result_type for all arrays/dtypes
     from ._creation_functions import empty