TYP: type annotations (#135)

* TYP: type annotations

* Python 3.9 fixes

* self-review

* code review

* Update array_api_strict/_array_object.py

Co-authored-by: Joren Hammudoglu <[email protected]>

* Apply suggestions from code review

Co-authored-by: Joren Hammudoglu <[email protected]>

* fix

* code review

* fixes

* normalize order

* Fancy indexing in `__getitem__` signature

* verbose Python scalar types

---------

Co-authored-by: Joren Hammudoglu <[email protected]>

Author: crusaderky

array_api_strict/__init__.py

@@ -16,6 +16,8 @@
 
 """
 
+from types import ModuleType
+
 __all__ = []
 
 # Warning: __array_api_version__ could change globally with
@@ -325,12 +327,16 @@
     ArrayAPIStrictFlags,
 )
 
-__all__ += ['set_array_api_strict_flags', 'get_array_api_strict_flags', 'reset_array_api_strict_flags', 'ArrayAPIStrictFlags']
+__all__ += [
+    'set_array_api_strict_flags',
+    'get_array_api_strict_flags',
+    'reset_array_api_strict_flags',
+    'ArrayAPIStrictFlags',
+    '__version__',
+]
 
 try:
-    from . import _version
-    __version__ = _version.__version__
-    del _version
+    from ._version import __version__  # type: ignore[import-not-found,unused-ignore]
 except ImportError:
     __version__ = "unknown"
 
@@ -340,7 +346,7 @@
 # use __getattr__. Note that linalg and fft are dynamically added and removed
 # from __all__ in set_array_api_strict_flags.
 
-def __getattr__(name):
+def __getattr__(name: str) -> ModuleType:
     if name in ['linalg', 'fft']:
         if name in get_array_api_strict_flags()['enabled_extensions']:
             if name == 'linalg':

array_api_strict/_array_object.py

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

array_api_strict/_constants.py

@@ -1,38 +1,37 @@
 from __future__ import annotations
 
-from ._array_object import Array
-from ._creation_functions import _check_device
+from dataclasses import dataclass
+
+import numpy as np
+
+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 +51,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 +64,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 +75,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>`.
 
@@ -112,26 +111,25 @@ class finfo_object:
     max: float
     min: float
     smallest_normal: float
-    dtype: Dtype
+    dtype: DType
 
 
 @dataclass
 class iinfo_object:
     bits: int
     max: int
     min: int
-    dtype: Dtype
+    dtype: DType
 
 
-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)
+    np_type = type._array if isinstance(type, Array) else type._np_dtype
+    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 +138,33 @@ 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)
+    np_type = type._array if isinstance(type, Array) else type._np_dtype
+    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 +193,10 @@ 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: DType | Array | bool | int | float | complex,
+) -> DType:
     """
     Array API compatible wrapper for :py:func:`np.result_type <numpy.result_type>`.
 
@@ -219,15 +218,15 @@ def result_type(*arrays_and_dtypes: Union[Array, Dtype, int, float, complex, boo
         A.append(a)
 
     # remove python scalars
-    A = [a for a in A if not isinstance(a, (bool, int, float, complex))]
+    B = [a for a in A if not isinstance(a, (bool, int, float, complex))]
 
-    if len(A) == 0:
+    if len(B) == 0:
         raise ValueError("at least one array or dtype is required")
-    elif len(A) == 1:
-        result = A[0]
+    elif len(B) == 1:
+        result = B[0]
     else:
-        t = A[0]
-        for t2 in A[1:]:
+        t = B[0]
+        for t2 in B[1:]:
             t = _result_type(t, t2)
         result = t
 

array_api_strict/_creation_functions.py

@@ -1,19 +1,27 @@
+from __future__ import annotations
+
+import builtins
 import warnings
+from typing import Any, Final
 
 import numpy as np
+import numpy.typing as npt
 
 # Note: we wrap the NumPy dtype objects in a bare class, so that none of the
 # additional methods and behaviors of NumPy dtype objects are exposed.
 
-class _DType:
-    def __init__(self, np_dtype):
-        np_dtype = np.dtype(np_dtype)
-        self._np_dtype = np_dtype
 
-    def __repr__(self):
+class DType:
+    _np_dtype: Final[np.dtype[Any]]
+    __slots__ = ("_np_dtype", "__weakref__")
+
+    def __init__(self, np_dtype: npt.DTypeLike):
+        self._np_dtype = np.dtype(np_dtype)
+
+    def __repr__(self) -> str:
         return f"array_api_strict.{self._np_dtype.name}"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> builtins.bool:
         # See https://github.com/numpy/numpy/pull/25370/files#r1423259515.
         # Avoid the user error of array_api_strict.float32 == numpy.float32,
         # which gives False. Making == error is probably too egregious, so
@@ -26,33 +34,38 @@ def __eq__(self, other):
 a NumPy native dtype object, but you probably don't want to do this. \
 array_api_strict dtype objects compare unequal to their NumPy equivalents. \
 Such cross-library comparison is not supported by the standard.""",
-            stacklevel=2)
-        if not isinstance(other, _DType):
+                stacklevel=2,
+            )
+        if not isinstance(other, DType):
             return NotImplemented
         return self._np_dtype == other._np_dtype
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         # Note: this is not strictly required
         # (https://github.com/data-apis/array-api/issues/582), but makes the
         # dtype objects much easier to work with here and elsewhere if they
         # can be used as dict keys.
         return hash(self._np_dtype)
 
 
-int8 = _DType("int8")
-int16 = _DType("int16")
-int32 = _DType("int32")
-int64 = _DType("int64")
-uint8 = _DType("uint8")
-uint16 = _DType("uint16")
-uint32 = _DType("uint32")
-uint64 = _DType("uint64")
-float32 = _DType("float32")
-float64 = _DType("float64")
-complex64 = _DType("complex64")
-complex128 = _DType("complex128")
+def _np_dtype(dtype: DType | None) -> np.dtype[Any] | None:
+    return dtype._np_dtype if dtype is not None else None
+
+
+int8 = DType("int8")
+int16 = DType("int16")
+int32 = DType("int32")
+int64 = DType("int64")
+uint8 = DType("uint8")
+uint16 = DType("uint16")
+uint32 = DType("uint32")
+uint64 = DType("uint64")
+float32 = DType("float32")
+float64 = DType("float64")
+complex64 = DType("complex64")
+complex128 = DType("complex128")
 # Note: This name is changed
-bool = _DType("bool")
+bool = DType("bool")
 
 _all_dtypes = (
     int8,
@@ -212,7 +225,7 @@ def __hash__(self):
 }
 
 
-def _result_type(type1, type2):
+def _result_type(type1: DType, type2: DType) -> DType:
     if (type1, type2) in _promotion_table:
         return _promotion_table[type1, type2]
     raise TypeError(f"{type1} and {type2} cannot be type promoted together")

array_api_strict/_data_type_functions.py

@@ -1,51 +1,50 @@
 from __future__ import annotations
 
+import numpy as np
+
+from ._array_object import Array
+from ._creation_functions import asarray
+from ._data_type_functions import broadcast_to, iinfo
 from ._dtypes import (
     _boolean_dtypes,
-    _floating_dtypes,
-    _real_floating_dtypes,
     _complex_floating_dtypes,
+    _dtype_categories,
+    _floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
-    _real_numeric_dtypes,
     _numeric_dtypes,
+    _real_floating_dtypes,
+    _real_numeric_dtypes,
     _result_type,
-    _dtype_categories,
 )
-from ._array_object import Array
 from ._flags import requires_api_version
-from ._creation_functions import asarray
-from ._data_type_functions import broadcast_to, iinfo
 from ._helpers import _maybe_normalize_py_scalars
 
-from typing import Optional, Union
-
-import numpy as np
-
 
 def _binary_ufunc_proto(x1, x2, dtype_category, func_name, np_func):
     """Base implementation of a binary function, `func_name`, defined for
-       dtypes from `dtype_category`
+    dtypes from `dtype_category`
     """
     x1, x2 = _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name)
 
     if x1.device != x2.device:
-        raise ValueError(f"Arrays from two different devices ({x1.device} and {x2.device}) can not be combined.")
+        raise ValueError(
+            f"Arrays from two different devices ({x1.device} and {x2.device}) can not be combined."
+        )
     # Call result type here just to raise on disallowed type combinations
     _result_type(x1.dtype, x2.dtype)
     x1, x2 = Array._normalize_two_args(x1, x2)
     return Array._new(np_func(x1._array, x2._array), device=x1.device)
 
 
-_binary_docstring_template=\
-"""
+_binary_docstring_template = """
 Array API compatible wrapper for :py:func:`np.%s <numpy.%s>`.
 
 See its docstring for more information.
 """
 
 
-def create_binary_func(func_name, dtype_category, np_func):
+def _create_binary_func(func_name, dtype_category, np_func):
     def inner(x1, x2, /) -> Array:
         return _binary_ufunc_proto(x1, x2, dtype_category, func_name, np_func)
     return inner
@@ -58,7 +57,7 @@ def inner(x1, x2, /) -> Array:
     "real numeric": "int | float | Array",
     "numeric": "int | float | complex | Array",
     "integer": "int | Array",
-    "integer or boolean": "int | bool | Array",
+    "integer or boolean": "bool | int | Array",
     "boolean": "bool | Array",
     "real floating-point": "float | Array",
     "complex floating-point": "complex | Array",
@@ -75,7 +74,7 @@ def inner(x1, x2, /) -> Array:
     "bitwise_xor": "integer or boolean",
     "_bitwise_left_shift": "integer",  # leading underscore deliberate
     "_bitwise_right_shift": "integer",
-    # XXX: copysign: real fp or numeric? 
+    # XXX: copysign: real fp or numeric?
     "copysign": "real floating-point",
     "divide": "floating-point",
     "equal": "all",
@@ -105,7 +104,7 @@ def inner(x1, x2, /) -> Array:
     "atan2": "arctan2",
     "_bitwise_left_shift": "left_shift",
     "_bitwise_right_shift": "right_shift",
-    "pow": "power"
+    "pow": "power",
 }
 
 
@@ -117,7 +116,7 @@ def inner(x1, x2, /) -> Array:
     numpy_name = _numpy_renames.get(func_name, func_name)
     np_func = getattr(np, numpy_name)
 
-    func = create_binary_func(func_name, dtype_category, np_func)
+    func = _create_binary_func(func_name, dtype_category, np_func)
     func.__name__ = func_name
 
     func.__doc__ = _binary_docstring_template % (numpy_name, numpy_name)
@@ -153,7 +152,7 @@ def bitwise_right_shift(x1: int | Array, x2: int | Array, /) -> Array:
 
 
 # clean up to not pollute the namespace
-del func, create_binary_func
+del func, _create_binary_func
 
 
 def abs(x: Array, /) -> Array:
@@ -271,8 +270,8 @@ def ceil(x: Array, /) -> Array:
 def clip(
     x: Array,
     /,
-    min: Optional[Union[int, float, Array]] = None,
-    max: Optional[Union[int, float, Array]] = None,
+    min: Array | int | float | None = None,
+    max: Array | int | float | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.clip <numpy.clip>`.
@@ -351,6 +350,7 @@ def clip(
 
     def _isscalar(a):
         return isinstance(a, (int, float, type(None)))
+
     min_shape = () if _isscalar(min) else min.shape
     max_shape = () if _isscalar(max) else max.shape
 
@@ -584,6 +584,7 @@ def reciprocal(x: Array, /) -> Array:
         raise TypeError("Only floating-point dtypes are allowed in reciprocal")
     return Array._new(np.reciprocal(x._array), device=x.device)
 
+
 def round(x: Array, /) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.round <numpy.round>`.

array_api_strict/_dtypes.py

@@ -1,31 +1,29 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from collections.abc import Sequence
+from typing import Literal
 
-if TYPE_CHECKING:
-    from typing import Union, Optional, Literal
-    from ._typing import Device, Dtype as DType
-    from collections.abc import Sequence
+import numpy as np
 
+from ._array_object import ALL_DEVICES, Array, Device
+from ._data_type_functions import astype
 from ._dtypes import (
+    DType,
+    _complex_floating_dtypes,
     _floating_dtypes,
     _real_floating_dtypes,
-    _complex_floating_dtypes,
-    float32,
     complex64,
+    float32,
 )
-from ._array_object import Array, ALL_DEVICES
-from ._data_type_functions import astype
 from ._flags import requires_extension
 
-import numpy as np
 
 @requires_extension('fft')
 def fft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -48,7 +46,7 @@ def ifft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -71,8 +69,8 @@ def fftn(
     x: Array,
     /,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
     """
@@ -94,8 +92,8 @@ def ifftn(
     x: Array,
     /,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
     """
@@ -117,7 +115,7 @@ def rfft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -140,7 +138,7 @@ def irfft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -163,8 +161,8 @@ def rfftn(
     x: Array,
     /,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
     """
@@ -186,8 +184,8 @@ def irfftn(
     x: Array,
     /,
     *,
-    s: Sequence[int] = None,
-    axes: Sequence[int] = None,
+    s: Sequence[int] | None = None,
+    axes: Sequence[int] | None = None,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
     """
@@ -209,7 +207,7 @@ def hfft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -232,7 +230,7 @@ def ihfft(
     x: Array,
     /,
     *,
-    n: Optional[int] = None,
+    n: int | None = None,
     axis: int = -1,
     norm: Literal["backward", "ortho", "forward"] = "backward",
 ) -> Array:
@@ -256,8 +254,8 @@ def fftfreq(
     /,
     *,
     d: float = 1.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.fft.fftfreq <numpy.fft.fftfreq>`.
@@ -280,8 +278,8 @@ def rfftfreq(
     /,
     *,
     d: float = 1.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.fft.rfftfreq <numpy.fft.rfftfreq>`.
@@ -299,7 +297,7 @@ def rfftfreq(
     return Array._new(np_result, device=device)
 
 @requires_extension('fft')
-def fftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
+def fftshift(x: Array, /, *, axes: int | Sequence[int] | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.fftshift <numpy.fft.fftshift>`.
 
@@ -310,7 +308,7 @@ def fftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
     return Array._new(np.fft.fftshift(x._array, axes=axes), device=x.device)
 
 @requires_extension('fft')
-def ifftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
+def ifftshift(x: Array, /, *, axes: int | Sequence[int] | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.ifftshift <numpy.fft.ifftshift>`.
 

array_api_strict/_elementwise_functions.py

@@ -12,17 +12,32 @@
 
 """
 
+from __future__ import annotations
+
 import functools
 import os
 import warnings
+from collections.abc import Callable
+from types import TracebackType
+from typing import TYPE_CHECKING, Any, Collection, TypeVar, cast
 
 import array_api_strict
 
+if TYPE_CHECKING:
+    # TODO import from typing (requires Python >= 3.10)
+    from typing_extensions import ParamSpec
+
+    P = ParamSpec("P")
+
+T = TypeVar("T")
+_CallableT = TypeVar("_CallableT", bound=Callable[..., object])
+
+
 supported_versions = (
     "2021.12",
     "2022.12",
     "2023.12",
-    "2024.12"
+    "2024.12",
 )
 
 draft_version = "2025.12"
@@ -43,19 +58,23 @@
     "fft": "2022.12",
 }
 
-ENABLED_EXTENSIONS = default_extensions = (
+default_extensions: tuple[str, ...] = (
     "linalg",
     "fft",
 )
+ENABLED_EXTENSIONS = default_extensions
+
+
 # Public functions
 
+
 def set_array_api_strict_flags(
     *,
-    api_version=None,
-    boolean_indexing=None,
-    data_dependent_shapes=None,
-    enabled_extensions=None,
-):
+    api_version: str | None = None,
+    boolean_indexing: bool | None = None,
+    data_dependent_shapes: bool | None = None,
+    enabled_extensions: Collection[str] | None = None,
+) -> None:
     """
     Set the array-api-strict flags to the specified values.
 
@@ -178,7 +197,8 @@ def set_array_api_strict_flags(
         draft_version=draft_version,
     )
 
-def get_array_api_strict_flags():
+
+def get_array_api_strict_flags() -> dict[str, Any]:
     """
     Get the current array-api-strict flags.
 
@@ -228,7 +248,7 @@ def get_array_api_strict_flags():
     }
 
 
-def reset_array_api_strict_flags():
+def reset_array_api_strict_flags() -> None:
     """
     Reset the array-api-strict flags to their default values.
 
@@ -300,8 +320,19 @@ class ArrayAPIStrictFlags:
     reset_array_api_strict_flags: Reset the flags to their default values.
 
     """
-    def __init__(self, *, api_version=None, boolean_indexing=None,
-                 data_dependent_shapes=None, enabled_extensions=None):
+
+    kwargs: dict[str, Any]
+    old_flags: dict[str, Any]
+    __slots__ = ("kwargs", "old_flags")
+
+    def __init__(
+        self,
+        *,
+        api_version: str | None = None,
+        boolean_indexing: bool | None = None,
+        data_dependent_shapes: bool | None = None,
+        enabled_extensions: Collection[str] | None = None,
+    ):
         self.kwargs = {
             "api_version": api_version,
             "boolean_indexing": boolean_indexing,
@@ -310,12 +341,19 @@ def __init__(self, *, api_version=None, boolean_indexing=None,
         }
         self.old_flags = get_array_api_strict_flags()
 
-    def __enter__(self):
+    def __enter__(self) -> None:
         set_array_api_strict_flags(**self.kwargs)
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+        /,
+    ) -> None:
         set_array_api_strict_flags(**self.old_flags)
 
+
 # Private functions
 
 ENVIRONMENT_VARIABLES = [
@@ -325,8 +363,9 @@ def __exit__(self, exc_type, exc_value, traceback):
     "ARRAY_API_STRICT_ENABLED_EXTENSIONS",
 ]
 
-def set_flags_from_environment():
-    kwargs = {}
+
+def set_flags_from_environment() -> None:
+    kwargs: dict[str, Any] = {}
     if "ARRAY_API_STRICT_API_VERSION" in os.environ:
         kwargs["api_version"] = os.environ["ARRAY_API_STRICT_API_VERSION"]
 
@@ -346,41 +385,49 @@ def set_flags_from_environment():
     # linalg and fft to __all__
     set_array_api_strict_flags(**kwargs)
 
+
 set_flags_from_environment()
 
 # Decorators
 
-def requires_api_version(version):
-    def decorator(func):
+
+def requires_api_version(version: str) -> Callable[[_CallableT], _CallableT]:
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
         @functools.wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             if version > API_VERSION:
                 raise RuntimeError(
                     f"The function {func.__name__} requires API version {version} or later, "
                     f"but the current API version for array-api-strict is {API_VERSION}"
                 )
             return func(*args, **kwargs)
+
         return wrapper
-    return decorator
 
-def requires_data_dependent_shapes(func):
+    return cast(Callable[[_CallableT], _CallableT], decorator)
+
+
+def requires_data_dependent_shapes(func: Callable[P, T]) -> Callable[P, T]:
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
         if not DATA_DEPENDENT_SHAPES:
             raise RuntimeError(f"The function {func.__name__} requires data-dependent shapes, but the data_dependent_shapes flag has been disabled for array-api-strict")
         return func(*args, **kwargs)
     return wrapper
 
-def requires_extension(extension):
-    def decorator(func):
+
+def requires_extension(extension: str) -> Callable[[_CallableT], _CallableT]:
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
         @functools.wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             if extension not in ENABLED_EXTENSIONS:
                 if extension == 'linalg' \
                    and func.__name__ in ['matmul', 'tensordot',
                                          'matrix_transpose', 'vecdot']:
                     raise RuntimeError(f"The linalg extension has been disabled for array-api-strict. However, {func.__name__} is also present in the main array_api_strict namespace and may be used from there.")
                 raise RuntimeError(f"The function {func.__name__} requires the {extension} extension, but it has been disabled for array-api-strict")
             return func(*args, **kwargs)
+
         return wrapper
-    return decorator
+
+    return cast(Callable[[_CallableT], _CallableT], decorator)

array_api_strict/_fft.py

@@ -1,18 +1,24 @@
-"""Private helper routines.
-"""
+"""Private helper routines."""
 
-from ._flags import get_array_api_strict_flags
+from __future__ import annotations
+
+from ._array_object import Array
 from ._dtypes import _dtype_categories
+from ._flags import get_array_api_strict_flags
 
 _py_scalars = (bool, int, float, complex)
 
 
-def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
-
+def _maybe_normalize_py_scalars(
+    x1: Array | bool | int | float | complex,
+    x2: Array | bool | int | float | complex,
+    dtype_category: str,
+    func_name: str,
+) -> tuple[Array, Array]:
     flags = get_array_api_strict_flags()
     if flags["api_version"] < "2024.12":
         # scalars will fail at the call site
-        return x1, x2
+        return x1, x2  # type: ignore[return-value]
 
     _allowed_dtypes = _dtype_categories[dtype_category]
 
@@ -34,4 +40,3 @@ def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
             raise TypeError(f"Only {dtype_category} dtypes are allowed in {func_name}(...). "
                             f"Got {x1.dtype} and {x2.dtype}.")
     return x1, x2
-

array_api_strict/_flags.py

@@ -1,17 +1,13 @@
 from __future__ import annotations
 
+import numpy as np
+
 from ._array_object import Array
 from ._dtypes import _integer_dtypes
 from ._flags import requires_api_version
 
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from typing import Optional
 
-import numpy as np
-
-def take(x: Array, indices: Array, /, *, axis: Optional[int] = None) -> Array:
+def take(x: Array, indices: Array, /, *, axis: int | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.take <numpy.take>`.
 
@@ -27,6 +23,7 @@ def take(x: Array, indices: Array, /, *, axis: Optional[int] = None) -> Array:
         raise ValueError(f"Arrays from two different devices ({x.device} and {indices.device}) can not be combined.")
     return Array._new(np.take(x._array, indices._array, axis=axis), device=x.device)
 
+
 @requires_api_version('2024.12')
 def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
     """

array_api_strict/_helpers.py

@@ -1,25 +1,22 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 import numpy as np
 
-if TYPE_CHECKING:
-    from typing import Optional, Union, Tuple, List
-    from ._typing import device, DefaultDataTypes, DataTypes, Capabilities
-
-from ._array_object import ALL_DEVICES, CPU_DEVICE
+from . import _dtypes as dt
+from ._array_object import ALL_DEVICES, CPU_DEVICE, Device
 from ._flags import get_array_api_strict_flags, requires_api_version
-from ._dtypes import bool, int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64, complex64, complex128
+from ._typing import Capabilities, DataTypes, DefaultDataTypes
+
 
 @requires_api_version('2023.12')
 class __array_namespace_info__:
     @requires_api_version('2023.12')
     def capabilities(self) -> Capabilities:
         flags = get_array_api_strict_flags()
-        res = {"boolean indexing": flags['boolean_indexing'],
-                "data-dependent shapes": flags['data_dependent_shapes'],
-                }
+        res: Capabilities = {  # type: ignore[typeddict-item]
+            "boolean indexing": flags['boolean_indexing'],
+            "data-dependent shapes": flags['data_dependent_shapes'],
+        }
         if flags['api_version'] >= '2024.12':
             # maxdims is 32 for NumPy 1.x and 64 for NumPy 2.0. Eventually we will
             # drop support for NumPy 1 but for now, just compute the number
@@ -36,104 +33,104 @@ def capabilities(self) -> Capabilities:
         return res
 
     @requires_api_version('2023.12')
-    def default_device(self) -> device:
+    def default_device(self) -> Device:
         return CPU_DEVICE
 
     @requires_api_version('2023.12')
     def default_dtypes(
         self,
         *,
-        device: Optional[device] = None,
+        device: Device | None = None,
     ) -> DefaultDataTypes:
         return {
-            "real floating": float64,
-            "complex floating": complex128,
-            "integral": int64,
-            "indexing": int64,
+            "real floating": dt.float64,
+            "complex floating": dt.complex128,
+            "integral": dt.int64,
+            "indexing": dt.int64,
         }
 
     @requires_api_version('2023.12')
     def dtypes(
         self,
         *,
-        device: Optional[device] = None,
-        kind: Optional[Union[str, Tuple[str, ...]]] = None,
+        device: Device | None = None,
+        kind: str | tuple[str, ...] | None = None,
     ) -> DataTypes:
         if kind is None:
             return {
-                "bool": bool,
-                "int8": int8,
-                "int16": int16,
-                "int32": int32,
-                "int64": int64,
-                "uint8": uint8,
-                "uint16": uint16,
-                "uint32": uint32,
-                "uint64": uint64,
-                "float32": float32,
-                "float64": float64,
-                "complex64": complex64,
-                "complex128": complex128,
+                "bool": dt.bool,
+                "int8": dt.int8,
+                "int16": dt.int16,
+                "int32": dt.int32,
+                "int64": dt.int64,
+                "uint8": dt.uint8,
+                "uint16": dt.uint16,
+                "uint32": dt.uint32,
+                "uint64": dt.uint64,
+                "float32": dt.float32,
+                "float64": dt.float64,
+                "complex64": dt.complex64,
+                "complex128": dt.complex128,
             }
         if kind == "bool":
-            return {"bool": bool}
+            return {"bool": dt.bool}
         if kind == "signed integer":
             return {
-                "int8": int8,
-                "int16": int16,
-                "int32": int32,
-                "int64": int64,
+                "int8": dt.int8,
+                "int16": dt.int16,
+                "int32": dt.int32,
+                "int64": dt.int64,
             }
         if kind == "unsigned integer":
             return {
-                "uint8": uint8,
-                "uint16": uint16,
-                "uint32": uint32,
-                "uint64": uint64,
+                "uint8": dt.uint8,
+                "uint16": dt.uint16,
+                "uint32": dt.uint32,
+                "uint64": dt.uint64,
             }
         if kind == "integral":
             return {
-                "int8": int8,
-                "int16": int16,
-                "int32": int32,
-                "int64": int64,
-                "uint8": uint8,
-                "uint16": uint16,
-                "uint32": uint32,
-                "uint64": uint64,
+                "int8": dt.int8,
+                "int16": dt.int16,
+                "int32": dt.int32,
+                "int64": dt.int64,
+                "uint8": dt.uint8,
+                "uint16": dt.uint16,
+                "uint32": dt.uint32,
+                "uint64": dt.uint64,
             }
         if kind == "real floating":
             return {
-                "float32": float32,
-                "float64": float64,
+                "float32": dt.float32,
+                "float64": dt.float64,
             }
         if kind == "complex floating":
             return {
-                "complex64": complex64,
-                "complex128": complex128,
+                "complex64": dt.complex64,
+                "complex128": dt.complex128,
             }
         if kind == "numeric":
             return {
-                "int8": int8,
-                "int16": int16,
-                "int32": int32,
-                "int64": int64,
-                "uint8": uint8,
-                "uint16": uint16,
-                "uint32": uint32,
-                "uint64": uint64,
-                "float32": float32,
-                "float64": float64,
-                "complex64": complex64,
-                "complex128": complex128,
+                "int8": dt.int8,
+                "int16": dt.int16,
+                "int32": dt.int32,
+                "int64": dt.int64,
+                "uint8": dt.uint8,
+                "uint16": dt.uint16,
+                "uint32": dt.uint32,
+                "uint64": dt.uint64,
+                "float32": dt.float32,
+                "float64": dt.float64,
+                "complex64": dt.complex64,
+                "complex128": dt.complex128,
             }
         if isinstance(kind, tuple):
-            res = {}
+            res: DataTypes = {}
             for k in kind:
                 res.update(self.dtypes(kind=k))
             return res
         raise ValueError(f"unsupported kind: {kind!r}")
 
     @requires_api_version('2023.12')
-    def devices(self) -> List[device]:
+    def devices(self) -> list[Device]:
         return list(ALL_DEVICES)

array_api_strict/_indexing_functions.py

@@ -1,33 +1,25 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import partial
+from typing import Literal, NamedTuple
 
-from ._dtypes import (
-    _floating_dtypes,
-    _numeric_dtypes,
-    float32,
-    complex64,
-    complex128,
-)
+import numpy as np
+import numpy.linalg
+
+from ._array_object import Array
 from ._data_type_functions import finfo
-from ._manipulation_functions import reshape
+from ._dtypes import DType, _floating_dtypes, _numeric_dtypes, complex64, complex128
 from ._elementwise_functions import conj
-from ._array_object import Array
-from ._flags import requires_extension, get_array_api_strict_flags
+from ._flags import get_array_api_strict_flags, requires_extension
+from ._manipulation_functions import reshape
+from ._statistical_functions import _np_dtype_sumprod
 
 try:
-    from numpy._core.numeric import normalize_axis_tuple
+    from numpy._core.numeric import normalize_axis_tuple  # type: ignore[attr-defined]
 except ImportError:
-    from numpy.core.numeric import normalize_axis_tuple
+    from numpy.core.numeric import normalize_axis_tuple  # type: ignore[no-redef]
 
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from ._typing import Literal, Optional, Sequence, Tuple, Union, Dtype
-
-from typing import NamedTuple
-
-import numpy.linalg
-import numpy as np
 
 class EighResult(NamedTuple):
     eigenvalues: Array
@@ -175,7 +167,13 @@ def inv(x: Array, /) -> Array:
 # -np.inf, 'fro', 'nuc']]], but Literal does not support floating-point
 # literals.
 @requires_extension('linalg')
-def matrix_norm(x: Array, /, *, keepdims: bool = False, ord: Optional[Union[int, float, Literal['fro', 'nuc']]] = 'fro') -> Array:  # noqa: F821
+def matrix_norm(
+    x: Array,
+    /,
+    *,
+    keepdims: bool = False,
+    ord: float | Literal["fro", "nuc"] | None = "fro",
+) -> Array:  # noqa: F821
     """
     Array API compatible wrapper for :py:func:`np.linalg.norm <numpy.linalg.norm>`.
 
@@ -186,7 +184,10 @@ def matrix_norm(x: Array, /, *, keepdims: bool = False, ord: Optional[Union[int,
     if x.dtype not in _floating_dtypes:
         raise TypeError('Only floating-point dtypes are allowed in matrix_norm')
 
-    return Array._new(np.linalg.norm(x._array, axis=(-2, -1), keepdims=keepdims, ord=ord), device=x.device)
+    return Array._new(
+        np.linalg.norm(x._array, axis=(-2, -1), keepdims=keepdims, ord=ord),
+        device=x.device,
+    )
 
 
 @requires_extension('linalg')
@@ -206,7 +207,7 @@ def matrix_power(x: Array, n: int, /) -> Array:
 
 # Note: the keyword argument name rtol is different from np.linalg.matrix_rank
 @requires_extension('linalg')
-def matrix_rank(x: Array, /, *, rtol: Optional[Union[float, Array]] = None) -> Array:
+def matrix_rank(x: Array, /, *, rtol: float | Array | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.matrix_rank <numpy.matrix_rank>`.
 
@@ -218,13 +219,12 @@ def matrix_rank(x: Array, /, *, rtol: Optional[Union[float, Array]] = None) -> A
         raise np.linalg.LinAlgError("1-dimensional array given. Array must be at least two-dimensional")
     S = np.linalg.svd(x._array, compute_uv=False)
     if rtol is None:
-        tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * finfo(S.dtype).eps
+        tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * np.finfo(S.dtype).eps
     else:
-        if isinstance(rtol, Array):
-            rtol = rtol._array
+        rtol_np = rtol._array if isinstance(rtol, Array) else np.asarray(rtol)
         # Note: this is different from np.linalg.matrix_rank, which does not multiply
         # the tolerance by the largest singular value.
-        tol = S.max(axis=-1, keepdims=True)*np.asarray(rtol)[..., np.newaxis]
+        tol = S.max(axis=-1, keepdims=True) * rtol_np[..., np.newaxis]
     return Array._new(np.count_nonzero(S > tol, axis=-1), device=x.device)
 
 
@@ -252,7 +252,7 @@ def outer(x1: Array, x2: Array, /) -> Array:
 
 # Note: the keyword argument name rtol is different from np.linalg.pinv
 @requires_extension('linalg')
-def pinv(x: Array, /, *, rtol: Optional[Union[float, Array]] = None) -> Array:
+def pinv(x: Array, /, *, rtol: float | Array | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.linalg.pinv <numpy.linalg.pinv>`.
 
@@ -267,9 +267,8 @@ def pinv(x: Array, /, *, rtol: Optional[Union[float, Array]] = None) -> Array:
     # default tolerance by max(M, N).
     if rtol is None:
         rtol = max(x.shape[-2:]) * finfo(x.dtype).eps
-    if isinstance(rtol, Array):
-        rtol = rtol._array
-    return Array._new(np.linalg.pinv(x._array, rcond=rtol), device=x.device)
+    rtol_np = rtol._array if isinstance(rtol, Array) else rtol
+    return Array._new(np.linalg.pinv(x._array, rcond=rtol_np), device=x.device)
 
 @requires_extension('linalg')
 def qr(x: Array, /, *, mode: Literal['reduced', 'complete'] = 'reduced') -> QRResult:  # noqa: F821
@@ -312,14 +311,14 @@ def slogdet(x: Array, /) -> SlogdetResult:
 
 # To workaround this, the below is the code from np.linalg.solve except
 # only calling solve1 in the exactly 1D case.
-def _solve(a, b):
+def _solve(a: np.ndarray, b: np.ndarray) -> np.ndarray:
     try:
-        from numpy.linalg._linalg import (
+        from numpy.linalg._linalg import (  # type: ignore[attr-defined]
         _makearray, _assert_stacked_2d, _assert_stacked_square,
         _commonType, isComplexType, _raise_linalgerror_singular
         )
     except ImportError:
-        from numpy.linalg.linalg import (
+        from numpy.linalg.linalg import (  # type: ignore[attr-defined]
         _makearray, _assert_stacked_2d, _assert_stacked_square,
         _commonType, isComplexType, _raise_linalgerror_singular
         )
@@ -382,14 +381,14 @@ def svd(x: Array, /, *, full_matrices: bool = True) -> SVDResult:
 # Note: svdvals is not in NumPy (but it is in SciPy). It is equivalent to
 # np.linalg.svd(compute_uv=False).
 @requires_extension('linalg')
-def svdvals(x: Array, /) -> Union[Array, Tuple[Array, ...]]:
+def svdvals(x: Array, /) -> Array:
     if x.dtype not in _floating_dtypes:
         raise TypeError('Only floating-point dtypes are allowed in svdvals')
     return Array._new(np.linalg.svd(x._array, compute_uv=False), device=x.device)
 
 # Note: trace is the numpy top-level namespace, not np.linalg
 @requires_extension('linalg')
-def trace(x: Array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> Array:
+def trace(x: Array, /, *, offset: int = 0, dtype: DType | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.trace <numpy.trace>`.
 
@@ -398,27 +397,28 @@ def trace(x: Array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> Arr
     if x.dtype not in _numeric_dtypes:
         raise TypeError('Only numeric dtypes are allowed in trace')
 
-    # Note: trace() works the same as sum() and prod() (see
-    # _statistical_functions.py)
-    if dtype is None:
-        if get_array_api_strict_flags()['api_version'] < '2023.12':
-            if x.dtype == float32:
-                dtype = np.float64
-            elif x.dtype == complex64:
-                dtype = np.complex128
-    else:
-        dtype = dtype._np_dtype
+    # Note: trace() works the same as sum() and prod() (see _statistical_functions.py)
+    np_dtype = _np_dtype_sumprod(x, dtype)
+
     # Note: trace always operates on the last two axes, whereas np.trace
     # operates on the first two axes by default
-    return Array._new(np.asarray(np.trace(x._array, offset=offset, axis1=-2, axis2=-1, dtype=dtype)), device=x.device)
+    res = np.trace(x._array, offset=offset, axis1=-2, axis2=-1, dtype=np_dtype)
+    return Array._new(np.asarray(res), device=x.device)
 
 # Note: the name here is different from norm(). The array API norm is split
 # into matrix_norm and vector_norm().
 
 # The type for ord should be Optional[Union[int, float, Literal[np.inf,
 # -np.inf]]] but Literal does not support floating-point literals.
 @requires_extension('linalg')
-def vector_norm(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False, ord: Optional[Union[int, float]] = 2) -> Array:
+def vector_norm(
+    x: Array,
+    /,
+    *,
+    axis: int | tuple[int, ...] | None = None,
+    keepdims: bool = False,
+    ord: int | float = 2,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.linalg.norm <numpy.linalg.norm>`.
 
@@ -456,8 +456,8 @@ def vector_norm(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = No
         # We can't reuse np.linalg.norm(keepdims) because of the reshape hacks
         # above to avoid matrix norm logic.
         shape = list(x.shape)
-        _axis = normalize_axis_tuple(range(x.ndim) if axis is None else axis, x.ndim)
-        for i in _axis:
+        axis_tup = normalize_axis_tuple(range(x.ndim) if axis is None else axis, x.ndim)
+        for i in axis_tup:
             shape[i] = 1
         res = reshape(res, tuple(shape))
 
@@ -480,7 +480,13 @@ def matmul(x1: Array, x2: Array, /) -> Array:
 
 # Note: tensordot is the numpy top-level namespace but not in np.linalg
 @requires_extension('linalg')
-def tensordot(x1: Array, x2: Array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> Array:
+def tensordot(
+    x1: Array,
+    x2: Array,
+    /,
+    *,
+    axes: int | tuple[Sequence[int], Sequence[int]] = 2,
+) -> Array:
     from ._linear_algebra_functions import tensordot
     return tensordot(x1, x2, axes=axes)
 

array_api_strict/_info.py

@@ -7,16 +7,15 @@
 
 from __future__ import annotations
 
-from ._dtypes import _numeric_dtypes
+from collections.abc import Sequence
+
+import numpy as np
+import numpy.linalg
+
 from ._array_object import Array
+from ._dtypes import _numeric_dtypes
 from ._flags import get_array_api_strict_flags
 
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from ._typing import Sequence, Tuple, Union
-
-import numpy.linalg
-import numpy as np
 
 # Note: matmul is the numpy top-level namespace but not in np.linalg
 def matmul(x1: Array, x2: Array, /) -> Array:
@@ -38,7 +37,13 @@ def matmul(x1: Array, x2: Array, /) -> Array:
 # Note: tensordot is the numpy top-level namespace but not in np.linalg
 
 # Note: axes must be a tuple, unlike np.tensordot where it can be an array or array-like.
-def tensordot(x1: Array, x2: Array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> Array:
+def tensordot(
+    x1: Array,
+    x2: Array,
+    /,
+    *,
+    axes: int | tuple[Sequence[int], Sequence[int]] = 2,
+) -> Array:
     # Note: the restriction to numeric dtypes only is different from
     # np.tensordot.
     if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:

array_api_strict/_linalg.py

@@ -1,21 +1,17 @@
 from __future__ import annotations
 
+import numpy as np
+
 from ._array_object import Array
 from ._creation_functions import asarray
 from ._data_type_functions import astype, result_type
 from ._dtypes import _integer_dtypes, int64, uint64
-from ._flags import requires_api_version, get_array_api_strict_flags
-
-from typing import TYPE_CHECKING
+from ._flags import get_array_api_strict_flags, requires_api_version
 
-if TYPE_CHECKING:
-    from typing import List, Optional, Tuple, Union
-
-import numpy as np
 
 # Note: the function name is different here
 def concat(
-    arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: Optional[int] = 0
+    arrays: tuple[Array, ...] | list[Array], /, *, axis: int | None = 0
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.concatenate <numpy.concatenate>`.
@@ -29,8 +25,11 @@ def concat(
         raise ValueError("concat inputs must all be on the same device")
     result_device = arrays[0].device
 
-    arrays = tuple(a._array for a in arrays)
-    return Array._new(np.concatenate(arrays, axis=axis, dtype=dtype._np_dtype), device=result_device)
+    np_arrays = tuple(a._array for a in arrays)
+    return Array._new(
+        np.concatenate(np_arrays, axis=axis, dtype=dtype._np_dtype),
+        device=result_device,
+    )
 
 
 def expand_dims(x: Array, /, *, axis: int) -> Array:
@@ -42,7 +41,7 @@ def expand_dims(x: Array, /, *, axis: int) -> Array:
     return Array._new(np.expand_dims(x._array, axis), device=x.device)
 
 
-def flip(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> Array:
+def flip(x: Array, /, *, axis: int | tuple[int, ...] | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.flip <numpy.flip>`.
 
@@ -53,8 +52,8 @@ def flip(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) ->
 @requires_api_version('2023.12')
 def moveaxis(
     x: Array,
-    source: Union[int, Tuple[int, ...]],
-    destination: Union[int, Tuple[int, ...]],
+    source: int | tuple[int, ...],
+    destination: int | tuple[int, ...],
     /,
 ) -> Array:
     """
@@ -66,7 +65,7 @@ def moveaxis(
 
 # Note: The function name is different here (see also matrix_transpose).
 # Unlike transpose(), the axes argument is required.
-def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array:
+def permute_dims(x: Array, /, axes: tuple[int, ...]) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.transpose <numpy.transpose>`.
 
@@ -77,10 +76,10 @@ def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array:
 @requires_api_version('2023.12')
 def repeat(
     x: Array,
-    repeats: Union[int, Array],
+    repeats: int | Array,
     /,
     *,
-    axis: Optional[int] = None,
+    axis: int | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.repeat <numpy.repeat>`.
@@ -108,12 +107,9 @@ def repeat(
         repeats = astype(repeats, int64)
     return Array._new(np.repeat(x._array, repeats._array, axis=axis), device=x.device)
 
+
 # Note: the optional argument is called 'shape', not 'newshape'
-def reshape(x: Array,
-            /,
-            shape: Tuple[int, ...],
-            *,
-            copy: Optional[bool] = None) -> Array:
+def reshape(x: Array, /, shape: tuple[int, ...], *, copy: bool | None = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.reshape <numpy.reshape>`.
 
@@ -135,9 +131,9 @@ def reshape(x: Array,
 def roll(
     x: Array,
     /,
-    shift: Union[int, Tuple[int, ...]],
+    shift: int | tuple[int, ...],
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.roll <numpy.roll>`.
@@ -147,7 +143,7 @@ def roll(
     return Array._new(np.roll(x._array, shift, axis=axis), device=x.device)
 
 
-def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array:
+def squeeze(x: Array, /, axis: int | tuple[int, ...]) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.squeeze <numpy.squeeze>`.
 
@@ -161,7 +157,7 @@ def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array:
     return Array._new(np.squeeze(x._array, axis=axis), device=x.device)
 
 
-def stack(arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: int = 0) -> Array:
+def stack(arrays: tuple[Array, ...] | list[Array], /, *, axis: int = 0) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.stack <numpy.stack>`.
 
@@ -172,12 +168,12 @@ def stack(arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: int = 0) ->
     if len({a.device for a in arrays}) > 1:
         raise ValueError("concat inputs must all be on the same device")
     result_device = arrays[0].device
-    arrays = tuple(a._array for a in arrays)
-    return Array._new(np.stack(arrays, axis=axis), device=result_device)
+    np_arrays = tuple(a._array for a in arrays)
+    return Array._new(np.stack(np_arrays, axis=axis), device=result_device)
 
 
 @requires_api_version('2023.12')
-def tile(x: Array, repetitions: Tuple[int, ...], /) -> Array:
+def tile(x: Array, repetitions: tuple[int, ...], /) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.tile <numpy.tile>`.
 
@@ -190,7 +186,7 @@ def tile(x: Array, repetitions: Tuple[int, ...], /) -> Array:
 
 # Note: this function is new
 @requires_api_version('2023.12')
-def unstack(x: Array, /, *, axis: int = 0) -> Tuple[Array, ...]:
+def unstack(x: Array, /, *, axis: int = 0) -> tuple[Array, ...]:
     if not (-x.ndim <= axis < x.ndim):
         raise ValueError("axis out of range")
 

array_api_strict/_linear_algebra_functions.py

@@ -1,18 +1,17 @@
 from __future__ import annotations
 
-from ._array_object import Array
-from ._dtypes import _result_type, _real_numeric_dtypes, bool as _bool
-from ._flags import requires_data_dependent_shapes, requires_api_version, get_array_api_strict_flags
-from ._helpers import _maybe_normalize_py_scalars
-
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from typing import Literal, Optional, Tuple, Union
+from typing import Literal
 
 import numpy as np
 
+from ._array_object import Array
+from ._dtypes import _real_numeric_dtypes, _result_type
+from ._dtypes import bool as _bool
+from ._flags import requires_api_version, requires_data_dependent_shapes
+from ._helpers import _maybe_normalize_py_scalars
+
 
-def argmax(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -> Array:
+def argmax(x: Array, /, *, axis: int | None = None, keepdims: bool = False) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.argmax <numpy.argmax>`.
 
@@ -23,7 +22,7 @@ def argmax(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -
     return Array._new(np.asarray(np.argmax(x._array, axis=axis, keepdims=keepdims)), device=x.device)
 
 
-def argmin(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -> Array:
+def argmin(x: Array, /, *, axis: int | None = None, keepdims: bool = False) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.argmin <numpy.argmin>`.
 
@@ -35,7 +34,7 @@ def argmin(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -
 
 
 @requires_data_dependent_shapes
-def nonzero(x: Array, /) -> Tuple[Array, ...]:
+def nonzero(x: Array, /) -> tuple[Array, ...]:
     """
     Array API compatible wrapper for :py:func:`np.nonzero <numpy.nonzero>`.
 
@@ -52,7 +51,7 @@ def count_nonzero(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     """
@@ -71,7 +70,7 @@ def searchsorted(
     /,
     *,
     side: Literal["left", "right"] = "left",
-    sorter: Optional[Array] = None,
+    sorter: Array | None = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.searchsorted <numpy.searchsorted>`.
@@ -84,25 +83,29 @@ def searchsorted(
     if x1.device != x2.device:
         raise ValueError(f"Arrays from two different devices ({x1.device} and {x2.device}) can not be combined.")
 
-    sorter = sorter._array if sorter is not None else None
+    np_sorter = sorter._array if sorter is not None else None
     # TODO: The sort order of nans and signed zeros is implementation
     # dependent. Should we error/warn if they are present?
 
     # x1 must be 1-D, but NumPy already requires this.
-    return Array._new(np.searchsorted(x1._array, x2._array, side=side, sorter=sorter), device=x1.device)
+    return Array._new(
+        np.searchsorted(x1._array, x2._array, side=side, sorter=np_sorter), 
+        device=x1.device,
+    )
+
 
 def where(
     condition: Array,
-    x1: bool | int | float | complex | Array,
-    x2: bool | int | float | complex | Array, /
+    x1: Array | bool | int | float | complex,
+    x2: Array | bool | int | float | complex,
+    /,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.where <numpy.where>`.
 
     See its docstring for more information.
     """
-    if get_array_api_strict_flags()['api_version'] > '2023.12':
-        x1, x2 = _maybe_normalize_py_scalars(x1, x2, "all", "where")
+    x1, x2 = _maybe_normalize_py_scalars(x1, x2, "all", "where")
 
     # Call result type here just to raise on disallowed type combinations
     _result_type(x1.dtype, x2.dtype)

array_api_strict/_manipulation_functions.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
-from ._array_object import Array
-
-from ._flags import requires_data_dependent_shapes
-
 from typing import NamedTuple
 
 import numpy as np
 
+from ._array_object import Array
+from ._flags import requires_data_dependent_shapes
+
 # Note: np.unique() is split into four functions in the array API:
 # unique_all, unique_counts, unique_inverse, and unique_values (this is done
 # to remove polymorphic return types).
@@ -20,6 +19,7 @@
 # Note: The functions here return a namedtuple (np.unique() returns a normal
 # tuple).
 
+
 class UniqueAllResult(NamedTuple):
     values: Array
     indices: Array

array_api_strict/_searching_functions.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
-from ._array_object import Array
-from ._dtypes import _real_numeric_dtypes
+from typing import Literal
 
 import numpy as np
 
+from ._array_object import Array
+from ._dtypes import _real_numeric_dtypes
+
 
 # Note: the descending keyword argument is new in this function
 def argsort(
@@ -18,7 +20,7 @@ def argsort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
+    kind: Literal["stable", "quicksort"] = "stable" if stable else "quicksort"
     if not descending:
         res = np.argsort(x._array, axis=axis, kind=kind)
     else:
@@ -35,6 +37,7 @@ def argsort(
         res = max_i - res
     return Array._new(res, device=x.device)
 
+
 # Note: the descending keyword argument is new in this function
 def sort(
     x: Array, /, *, axis: int = -1, descending: bool = False, stable: bool = True
@@ -47,8 +50,7 @@ def sort(
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in sort")
     # Note: this keyword argument is different, and the default is different.
-    kind = "stable" if stable else "quicksort"
-    res = np.sort(x._array, axis=axis, kind=kind)
+    res = np.sort(x._array, axis=axis, kind="stable" if stable else "quicksort")
     if descending:
         res = np.flip(res, axis=axis)
     return Array._new(res, device=x.device)

array_api_strict/_set_functions.py

@@ -1,38 +1,36 @@
 from __future__ import annotations
 
+from typing import Any
+
+import numpy as np
+
+from ._array_object import Array
+from ._creation_functions import ones, zeros
 from ._dtypes import (
-    _real_floating_dtypes,
-    _real_numeric_dtypes,
+    DType,
     _floating_dtypes,
+    _np_dtype,
     _numeric_dtypes,
+    _real_floating_dtypes,
+    _real_numeric_dtypes,
+    complex64,
+    float32,
 )
-from ._array_object import Array
-from ._dtypes import float32, complex64
-from ._flags import requires_api_version, get_array_api_strict_flags
-from ._creation_functions import zeros, ones
+from ._flags import get_array_api_strict_flags, requires_api_version
 from ._manipulation_functions import concat
 
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union
-    from ._typing import Dtype
-
-import numpy as np
 
 @requires_api_version('2023.12')
 def cumulative_sum(
     x: Array,
     /,
     *,
-    axis: Optional[int] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | None = None,
+    dtype: DType | None = None,
     include_initial: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in cumulative_sum")
-    if dtype is not None:
-        dtype = dtype._np_dtype
 
     # TODO: The standard is not clear about what should happen when x.ndim == 0.
     if axis is None:
@@ -44,26 +42,23 @@ def cumulative_sum(
         if axis < 0:
             axis += x.ndim
         x = concat([zeros(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype), x], axis=axis)
-    return Array._new(np.cumsum(x._array, axis=axis, dtype=dtype), device=x.device)
+    return Array._new(np.cumsum(x._array, axis=axis, dtype=_np_dtype(dtype)), device=x.device)
 
 
 @requires_api_version('2024.12')
 def cumulative_prod(
     x: Array,
     /,
     *,
-    axis: Optional[int] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | None = None,
+    dtype: DType | None = None,
     include_initial: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in cumulative_prod")
     if x.ndim == 0:
         raise ValueError("Only ndim >= 1 arrays are allowed in cumulative_prod")
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-
     if axis is None:
         if x.ndim > 1:
             raise ValueError("axis must be specified in cumulative_prod for more than one dimension")
@@ -74,14 +69,14 @@ def cumulative_prod(
         if axis < 0:
             axis += x.ndim
         x = concat([ones(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype), x], axis=axis)
-    return Array._new(np.cumprod(x._array, axis=axis, dtype=dtype), device=x.device)
+    return Array._new(np.cumprod(x._array, axis=axis, dtype=_np_dtype(dtype)), device=x.device)
 
 
 def max(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _real_numeric_dtypes:
@@ -93,14 +88,15 @@ def mean(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
 
-    if get_array_api_strict_flags()['api_version'] > '2023.12':
-        allowed_dtypes = _floating_dtypes
-    else:
-        allowed_dtypes = _real_floating_dtypes
+    allowed_dtypes = (
+        _floating_dtypes
+        if get_array_api_strict_flags()['api_version'] > '2023.12'
+        else _real_floating_dtypes
+    )
 
     if x.dtype not in allowed_dtypes:
         raise TypeError("Only floating-point dtypes are allowed in mean")
@@ -111,45 +107,51 @@ def min(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in min")
     return Array._new(np.min(x._array, axis=axis, keepdims=keepdims), device=x.device)
 
 
+def _np_dtype_sumprod(x: Array, dtype: DType | None) -> np.dtype[Any] | None:
+    """In versions prior to 2023.12, sum() and prod() upcast for all
+    dtypes when dtype=None. For 2023.12, the behavior is the same as in
+    NumPy (only upcast for integral dtypes).
+    """
+    if dtype is None and get_array_api_strict_flags()['api_version'] < '2023.12':
+        if x.dtype == float32:
+            return np.float64  # type: ignore[return-value]
+        elif x.dtype == complex64:
+            return np.complex128  # type: ignore[return-value]
+    return _np_dtype(dtype)
+
+
 def prod(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | tuple[int, ...] | None = None,
+    dtype: DType | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in prod")
 
-    if dtype is None:
-        # Note: In versions prior to 2023.12, sum() and prod() upcast for all
-        # dtypes when dtype=None. For 2023.12, the behavior is the same as in
-        # NumPy (only upcast for integral dtypes).
-        if get_array_api_strict_flags()['api_version'] < '2023.12':
-            if x.dtype == float32:
-                dtype = np.float64
-            elif x.dtype == complex64:
-                dtype = np.complex128
-    else:
-        dtype = dtype._np_dtype
-    return Array._new(np.prod(x._array, dtype=dtype, axis=axis, keepdims=keepdims), device=x.device)
+    np_dtype = _np_dtype_sumprod(x, dtype)
+    return Array._new(
+        np.prod(x._array, dtype=np_dtype, axis=axis, keepdims=keepdims),
+        device=x.device,
+    )
 
 
 def std(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    correction: Union[int, float] = 0.0,
+    axis: int | tuple[int, ...] | None = None,
+    correction: int | float = 0.0,
     keepdims: bool = False,
 ) -> Array:
     # Note: the keyword argument correction is different here
@@ -162,33 +164,26 @@ def sum(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | tuple[int, ...] | None = None,
+    dtype: DType | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in sum")
 
-    if dtype is None:
-        # Note: In versions prior to 2023.12, sum() and prod() upcast for all
-        # dtypes when dtype=None. For 2023.12, the behavior is the same as in
-        # NumPy (only upcast for integral dtypes).
-        if get_array_api_strict_flags()['api_version'] < '2023.12':
-            if x.dtype == float32:
-                dtype = np.float64
-            elif x.dtype == complex64:
-                dtype = np.complex128
-    else:
-        dtype = dtype._np_dtype
-    return Array._new(np.sum(x._array, axis=axis, dtype=dtype, keepdims=keepdims), device=x.device)
+    np_dtype = _np_dtype_sumprod(x, dtype)
+    return Array._new(
+        np.sum(x._array, axis=axis, dtype=np_dtype, keepdims=keepdims),
+        device=x.device,
+    )
 
 
 def var(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    correction: Union[int, float] = 0.0,
+    axis: int | tuple[int, ...] | None = None,
+    correction: int | float = 0.0,
     keepdims: bool = False,
 ) -> Array:
     # Note: the keyword argument correction is different here

array_api_strict/_sorting_functions.py

@@ -8,82 +8,62 @@
 
 from __future__ import annotations
 
-__all__ = [
-    "Array",
-    "Device",
-    "Dtype",
-    "SupportsDLPack",
-    "SupportsBufferProtocol",
-    "PyCapsule",
-]
-
 import sys
+from typing import Any, Protocol, TypedDict, TypeVar
 
-from typing import (
-    Any,
-    TypedDict,
-    TypeVar,
-    Protocol,
-)
-
-from ._array_object import Array, _device
-from ._dtypes import _DType
-from ._info import __array_namespace_info__
+from ._dtypes import DType
 
 _T_co = TypeVar("_T_co", covariant=True)
 
+
 class NestedSequence(Protocol[_T_co]):
     def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
     def __len__(self, /) -> int: ...
 
 
-Device = _device
-
-Dtype = _DType
-
-Info = __array_namespace_info__
-
 if sys.version_info >= (3, 12):
     from collections.abc import Buffer as SupportsBufferProtocol
 else:
     SupportsBufferProtocol = Any
 
 PyCapsule = Any
 
+
 class SupportsDLPack(Protocol):
     def __dlpack__(self, /, *, stream: None = ...) -> PyCapsule: ...
 
+
 Capabilities = TypedDict(
-    "Capabilities", {"boolean indexing": bool, "data-dependent shapes": bool,
-                     "max dimensions": int}
+    "Capabilities",
+    {
+        "boolean indexing": bool,
+        "data-dependent shapes": bool,
+        "max dimensions": int,
+    },
 )
 
 DefaultDataTypes = TypedDict(
     "DefaultDataTypes",
     {
-        "real floating": Dtype,
-        "complex floating": Dtype,
-        "integral": Dtype,
-        "indexing": Dtype,
+        "real floating": DType,
+        "complex floating": DType,
+        "integral": DType,
+        "indexing": DType,
     },
 )
 
-DataTypes = TypedDict(
-    "DataTypes",
-    {
-        "bool": Dtype,
-        "float32": Dtype,
-        "float64": Dtype,
-        "complex64": Dtype,
-        "complex128": Dtype,
-        "int8": Dtype,
-        "int16": Dtype,
-        "int32": Dtype,
-        "int64": Dtype,
-        "uint8": Dtype,
-        "uint16": Dtype,
-        "uint32": Dtype,
-        "uint64": Dtype,
-    },
-    total=False,
-)
+
+class DataTypes(TypedDict, total=False):
+    bool: DType
+    float32: DType
+    float64: DType
+    complex64: DType
+    complex128: DType
+    int8: DType
+    int16: DType
+    int32: DType
+    int64: DType
+    uint8: DType
+    uint16: DType
+    uint32: DType
+    uint64: DType

array_api_strict/_statistical_functions.py

@@ -1,21 +1,20 @@
 from __future__ import annotations
 
-from ._array_object import Array
-from ._flags import requires_api_version
-from ._dtypes import _numeric_dtypes
-
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union
+from typing import Any
 
 import numpy as np
+import numpy.typing as npt
+
+from ._array_object import Array
+from ._dtypes import _numeric_dtypes
+from ._flags import requires_api_version
 
 
 def all(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     """
@@ -30,7 +29,7 @@ def any(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     """
@@ -40,15 +39,16 @@ def any(
     """
     return Array._new(np.asarray(np.any(x._array, axis=axis, keepdims=keepdims)), device=x.device)
 
+
 @requires_api_version('2024.12')
 def diff(
     x: Array,
     /,
     *,
     axis: int = -1,
     n: int = 1,
-    prepend: Optional[Array] = None,
-    append: Optional[Array] = None,
+    prepend: Array | None = None,
+    append: Array | None = None,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in diff")
@@ -57,7 +57,7 @@ def diff(
     # currently specified.
 
     # NumPy does not support prepend=None or append=None
-    kwargs = dict(axis=axis, n=n)
+    kwargs: dict[str, int | npt.NDArray[Any]] = {"axis": axis, "n": n}
     if prepend is not None:
         if prepend.device != x.device:
             raise ValueError(f"Arrays from two different devices ({prepend.device} and {x.device}) can not be combined.")

array_api_strict/_typing.py

@@ -1,11 +1,9 @@
-from typing import Callable
-
 import pytest
 
 import array_api_strict as xp
 
 
-def p(func: Callable, *args, **kwargs):
+def p(func, *args, **kwargs):
     f_sig = ", ".join(
         [str(a) for a in args] + [f"{k}={v}" for k, v in kwargs.items()]
     )

array_api_strict/_utility_functions.py

@@ -31,3 +31,18 @@ Repository = "https://github.com/data-apis/array-api-strict"
 [tool.setuptools_scm]
 version_file = "array_api_strict/_version.py"
 
+[tool.mypy]
+disallow_incomplete_defs = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+show_error_codes = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+warn_unreachable = true
+strict_bytes = true
+local_partial_types = true
+
+[[tool.mypy.overrides]]
+module = ["*.tests.*"]
+disallow_untyped_defs = false