code review

Author: crusaderky

array_api_strict/_array_object.py

@@ -16,15 +16,16 @@
 from __future__ import annotations
 
 import operator
+import sys
 from collections.abc import Iterator
 from enum import IntEnum
 from types import ModuleType
-from typing import Any, Literal, SupportsIndex
+from typing import Any, Final, Literal, SupportsIndex
 
 import numpy as np
 import numpy.typing as npt
 
-from ._creation_functions import _undef, Undef, asarray
+from ._creation_functions import Undef, _undef, asarray
 from ._dtypes import (
     DType,
     _all_dtypes,
@@ -42,14 +43,16 @@
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
 from ._typing import PyCapsule
 
-try:
-    from types import EllipsisType  # Python >=3.10
-except ImportError:
+if sys.version_info >= (3, 10):
+    from types import EllipsisType
+elif TYPE_CHECKING:
+    from typing_extensions import EllipsisType
+else:
     EllipsisType = type(Ellipsis)
 
 
 class Device:
-    _device: str
+    _device: Final[str]
     __slots__ = ("_device", "__weakref__")
 
     def __init__(self, device: str = "CPU_DEVICE"):
@@ -101,7 +104,7 @@ class Array:
     # Use a custom constructor instead of __init__, as manually initializing
     # this class is not supported API.
     @classmethod
-    def _new(cls, x: npt.NDArray[Any], /, device: Device | None) -> Array:
+    def _new(cls, x: npt.NDArray[Any] | np.generic, /, device: Device | None) -> Array:
         """
         This is a private method for initializing the array API Array
         object.

array_api_strict/_constants.py

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

array_api_strict/_creation_functions.py

@@ -12,10 +12,12 @@
 from ._typing import NestedSequence, SupportsBufferProtocol, SupportsDLPack
 
 if TYPE_CHECKING:
+    # TODO import from typing (requires Python >=3.13)
+    from typing_extensions import TypeIs
+
     # Circular import
     from ._array_object import Array, Device
 
-
 class Undef(Enum):
     UNDEF = 0
 
@@ -44,7 +46,7 @@ def _check_valid_dtype(dtype: DType | None) -> None:
         raise ValueError(f"dtype must be one of the supported dtypes, got {dtype!r}")
 
 
-def _supports_buffer_protocol(obj: object) -> bool:
+def _supports_buffer_protocol(obj: object) -> TypeIs[SupportsBufferProtocol]:
     try:
         memoryview(obj)  # type: ignore[arg-type]
     except TypeError:

array_api_strict/_data_type_functions.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from ._array_object import Array, Device
-from ._creation_functions import _check_device, _undef, Undef
+from ._creation_functions import Undef, _check_device, _undef
 from ._dtypes import (
     DType,
     _all_dtypes,

array_api_strict/_dtypes.py

@@ -2,7 +2,7 @@
 
 import builtins
 import warnings
-from typing import Any
+from typing import Any, Final
 
 import numpy as np
 import numpy.typing as npt
@@ -12,7 +12,7 @@
 
 
 class DType:
-    _np_dtype: np.dtype[Any]
+    _np_dtype: Final[np.dtype[Any]]
     __slots__ = ("_np_dtype", "__weakref__")
 
     def __init__(self, np_dtype: npt.DTypeLike):

array_api_strict/_flags.py

@@ -19,7 +19,7 @@
 import warnings
 from collections.abc import Callable
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Collection, TypeVar
+from typing import TYPE_CHECKING, Any, Collection, TypeVar, cast
 
 import array_api_strict
 
@@ -30,6 +30,7 @@
     P = ParamSpec("P")
 
 T = TypeVar("T")
+_CallableT = TypeVar("_CallableT", bound=Callable[..., object])
 
 
 supported_versions = (
@@ -389,8 +390,7 @@ def set_flags_from_environment() -> None:
 
 # Decorators
 
-
-def requires_api_version(version: str) -> Callable[[Callable], Callable]:
+def requires_api_version(version: str) -> Callable[[_CallableT], _CallableT]:
     def decorator(func: Callable[P, T]) -> Callable[P, T]:
         @functools.wraps(func)
         def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
@@ -403,7 +403,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
 
         return wrapper
 
-    return decorator
+    return cast(Callable[[_CallableT], _CallableT], decorator)
 
 
 def requires_data_dependent_shapes(func: Callable[P, T]) -> Callable[P, T]:
@@ -415,7 +415,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
     return wrapper
 
 
-def requires_extension(extension: str) -> Callable[[Callable], Callable]:
+def requires_extension(extension: str) -> Callable[[_CallableT], _CallableT]:
     def decorator(func: Callable[P, T]) -> Callable[P, T]:
         @functools.wraps(func)
         def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
@@ -426,5 +426,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
                     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/_linalg.py

@@ -9,13 +9,7 @@
 
 from ._array_object import Array
 from ._data_type_functions import finfo
-from ._dtypes import (
-    DType,
-    _floating_dtypes,
-    _numeric_dtypes,
-    complex64,
-    complex128,
-)
+from ._dtypes import DType, _floating_dtypes, _numeric_dtypes, complex64, complex128
 from ._elementwise_functions import conj
 from ._flags import get_array_api_strict_flags, requires_extension
 from ._manipulation_functions import reshape

array_api_strict/_typing.py

@@ -52,22 +52,18 @@ def __dlpack__(self, /, *, stream: None = ...) -> PyCapsule: ...
     },
 )
 
-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

pyproject.toml

@@ -40,6 +40,8 @@ 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.*"]