code review

Author: crusaderky

array_api_compat/common/_aliases.py

@@ -177,7 +177,7 @@ class UniqueInverseResult(NamedTuple):
     inverse_indices: Array
 
 
-def _unique_kwargs(xp: Namespace) -> dict[str, Any]:
+def _unique_kwargs(xp: Namespace) -> dict[str, bool]:
     # Older versions of NumPy and CuPy do not have equal_nan. Rather than
     # trying to parse version numbers, just check if equal_nan is in the
     # signature.

array_api_compat/cupy/_aliases.py

@@ -126,10 +126,10 @@ def astype(
 
 # cupy.count_nonzero does not have keepdims
 def count_nonzero(
-    x: ndarray,
+    x: Array,
     axis=None,
     keepdims=False
-) -> ndarray:
+) -> Array:
    result = cp.count_nonzero(x, axis)
    if keepdims:
        if axis is None:

array_api_compat/cupy/_typing.py

@@ -3,31 +3,29 @@
 __all__ = ["Array", "DType", "Device"]
 _all_ignore = ["cp"]
 
-from typing import Union
+from typing import Union, TYPE_CHECKING
 
 import cupy as cp
 from cupy import ndarray as Array
 from cupy.cuda.device import Device
 
-try:
+if TYPE_CHECKING:
+    # NumPy 1.x on Python 3.9 and 3.10 fails to parse np.dtype[]
     DType = cp.dtype[
-        Union[
-            cp.intp,
-            cp.int8,
-            cp.int16,
-            cp.int32,
-            cp.int64,
-            cp.uint8,
-            cp.uint16,
-            cp.uint32,
-            cp.uint64,
-            cp.float32,
-            cp.float64,
-            cp.complex64,
-            cp.complex128,
-            cp.bool_,
-        ]
+        cp.intp
+        | cp.int8
+        | cp.int16
+        | cp.int32
+        | cp.int64
+        | cp.uint8
+        | cp.uint16
+        | cp.uint32
+        | cp.uint64
+        | cp.float32
+        | cp.float64
+        | cp.complex64
+        | cp.complex128
+        | cp.bool_
     ]
-except TypeError:
-    # NumPy 1.x on Python 3.9 and 3.10
+else:
     DType = cp.dtype

array_api_compat/numpy/_aliases.py

@@ -128,10 +128,10 @@ def astype(
 # count_nonzero returns a python int for axis=None and keepdims=False
 # https://github.com/numpy/numpy/issues/17562
 def count_nonzero(
-    x : ndarray,
+    x : Array,
     axis=None,
     keepdims=False
-) -> ndarray:
+) -> Array:
     result = np.count_nonzero(x, axis=axis, keepdims=keepdims)
     if axis is None and not keepdims:
         return np.asarray(result)

array_api_compat/numpy/_typing.py

@@ -3,31 +3,29 @@
 __all__ = ["Array", "DType", "Device"]
 _all_ignore = ["np"]
 
-from typing import Literal, Union
+from typing import Literal, TYPE_CHECKING
 
 import numpy as np
 from numpy import ndarray as Array
 
 Device = Literal["cpu"]
-try:
+if TYPE_CHECKING:
+    # NumPy 1.x on Python 3.9 and 3.10 fails to parse np.dtype[]
     DType = np.dtype[
-        Union[
-            np.intp,
-            np.int8,
-            np.int16,
-            np.int32,
-            np.int64,
-            np.uint8,
-            np.uint16,
-            np.uint32,
-            np.uint64,
-            np.float32,
-            np.float64,
-            np.complex64,
-            np.complex128,
-            np.bool_,
-        ]
+        np.intp
+        | np.int8
+        | np.int16
+        | np.int32
+        | np.int64
+        | np.uint8
+        | np.uint16
+        | np.uint32
+        | np.uint64
+        | np.float32
+        | np.float64
+        | np.complex64
+        | np.complex128
+        | np.bool
     ]
-except TypeError:
-    # NumPy 1.x on Python 3.9 and 3.10
+else:
     DType = np.dtype

array_api_compat/torch/_aliases.py

@@ -5,6 +5,7 @@
 from typing import List, Optional, Sequence, Tuple, Union
 
 import torch
+
 from .._internal import get_xp
 from ..common import _aliases
 from ._info import __array_namespace_info__

tests/test_all.py

@@ -17,14 +17,14 @@
 import pytest
 import typing
 
-TYPING_NAMES = {
+TYPING_NAMES = frozenset((
     "Array",
     "Device",
     "DType",
     "Namespace",
     "NestedSequence",
     "SupportsBufferProtocol",
-}
+))
 
 @pytest.mark.parametrize("library", ["common"] + wrapped_libraries)
 def test_all(library):