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

array_api_strict/_fft.py

@@ -1,31 +1,27 @@
-from __future__ import annotations
+from collections.abc import Sequence
+from typing import Literal
 
-from typing import TYPE_CHECKING
-
-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 +44,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 +67,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 +90,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 +113,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 +136,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 +159,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 +182,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 +205,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 +228,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 +252,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 +276,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 +295,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 +306,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/_flags.py

@@ -11,18 +11,25 @@ None of these functions are part of the standard itself. A typical array API
 library will only support one particular configuration of these flags.
 
 """
-
 import functools
 import os
 import warnings
+from collections.abc import Callable
+from types import TracebackType
+from typing import Any, Collection, ParamSpec, TypeVar, cast
 
 import array_api_strict
 
+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 +50,23 @@ extension_versions = {
     "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 +189,8 @@ if set_array_api_strict_flags.__doc__ is not None:
         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 +240,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 +312,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 +333,19 @@ class ArrayAPIStrictFlags:
         }
         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 +355,9 @@ ENVIRONMENT_VARIABLES = [
     "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,35 +377,41 @@ 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',
@@ -382,5 +419,7 @@ def requires_extension(extension):
                     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/_helpers.py

@@ -1,37 +1,46 @@
-"""Private helper routines.
-"""
+"""Private helper routines."""
 
-from ._flags import get_array_api_strict_flags
+from ._array_object import Array
 from ._dtypes import _dtype_categories
+from ._flags import get_array_api_strict_flags
 
-_py_scalars = (bool, int, float, complex)
-
+_PY_SCALARS = (bool, int, float, complex)
 
-def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
 
+def _maybe_normalize_py_scalars(
+    x1: Array | complex,
+    x2: Array | 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]
 
-    if isinstance(x1, _py_scalars):
-        if isinstance(x2, _py_scalars):
+    # Disallow subclasses, e.g. np.float64 and np.complex128
+    if type(x1) in _PY_SCALARS:
+        if type(x2) in _PY_SCALARS:
             raise TypeError(f"Two scalars not allowed, got {type(x1) = } and {type(x2) =}")
-        # x2 must be an array
+        if not isinstance(x2, Array):
+            raise TypeError(f"Argument is neither an Array nor a Python scalar: {type(x2)=} ")
         if x2.dtype not in _allowed_dtypes:
             raise TypeError(f"Only {dtype_category} dtypes are allowed {func_name}. Got {x2.dtype}.")
         x1 = x2._promote_scalar(x1)
 
-    elif isinstance(x2, _py_scalars):
-        # x1 must be an array
+    elif type(x2) in _PY_SCALARS:
+        if not isinstance(x1, Array):
+            raise TypeError(f"Argument is neither an Array nor a Python scalar: {type(x2)=} ")
         if x1.dtype not in _allowed_dtypes:
             raise TypeError(f"Only {dtype_category} dtypes are allowed {func_name}. Got {x1.dtype}.")
         x2 = x1._promote_scalar(x2)
     else:
+        if not isinstance(x1, Array) or not isinstance(x2, Array):
+            raise TypeError(f"Argument(s) are neither Array nor Python scalars: {type(x1)=} and {type(x2)=}")
+
         if x1.dtype not in _allowed_dtypes or x2.dtype not in _allowed_dtypes:
-            raise TypeError(f"Only {dtype_category} dtypes are allowed {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/_indexing_functions.py

@@ -1,17 +1,11 @@
-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 +21,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/_info.py

@@ -1,25 +1,20 @@
-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 +31,104 @@ class __array_namespace_info__:
         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)