address review

Author: lucascolley

src/array_api_extra/_lib/_backends.py

@@ -18,23 +18,28 @@ class Backend(Enum):  # numpydoc ignore=PR01,PR02  # type: ignore[no-subclass-an
     ----------
     value : str
         String describing the backend.
-    library_name : str
-        Name of the array library of the backend.
+    is_namespace : Callable[[ModuleType], bool]
+        Function to check whether an input module is the array namespace
+        corresponding to the backend.
     module_name : str
         Name of the backend's module.
     """
 
-    ARRAY_API_STRICT = "array_api_strict", "array_api_strict", "array_api_strict"
-    NUMPY = "numpy", "numpy", "numpy"
-    NUMPY_READONLY = "numpy_readonly", "numpy", "numpy"
-    CUPY = "cupy", "cupy", "cupy"
-    TORCH = "torch", "torch", "torch"
-    DASK_ARRAY = "dask.array", "dask", "dask.array"
-    SPARSE = "sparse", "pydata_sparse", "sparse"
-    JAX_NUMPY = "jax.numpy", "jax", "jax.numpy"
+    ARRAY_API_STRICT = (
+        "array_api_strict",
+        _compat.is_array_api_strict_namespace,
+        "array_api_strict",
+    )
+    NUMPY = "numpy", _compat.is_numpy_namespace, "numpy"
+    NUMPY_READONLY = "numpy_readonly", _compat.is_numpy_namespace, "numpy"
+    CUPY = "cupy", _compat.is_cupy_namespace, "cupy"
+    TORCH = "torch", _compat.is_torch_namespace, "torch"
+    DASK_ARRAY = "dask.array", _compat.is_dask_namespace, "dask.array"
+    SPARSE = "sparse", _compat.is_pydata_sparse_namespace, "sparse"
+    JAX_NUMPY = "jax.numpy", _compat.is_jax_namespace, "jax.numpy"
 
     def __new__(
-        cls, value: str, _library_name: str, _module_name: str
+        cls, value: str, _is_namespace: Callable[[ModuleType], bool], _module_name: str
     ):  # numpydoc ignore=GL08
         obj = object.__new__(cls)
         obj._value_ = value
@@ -43,30 +48,12 @@ def __new__(
     def __init__(
         self,
         value: str,  # noqa: ARG002  # pylint: disable=unused-argument
-        library_name: str,
+        is_namespace: Callable[[ModuleType], bool],
         module_name: str,
     ):  # numpydoc ignore=GL08
-        self.library_name = library_name
+        self.is_namespace = is_namespace
         self.module_name = module_name
 
     def __str__(self) -> str:  # type: ignore[explicit-override]  # pyright: ignore[reportImplicitOverride]  # numpydoc ignore=RT01
         """Pretty-print parameterized test names."""
         return cast(str, self.value)
-
-    def is_namespace(self, xp: ModuleType) -> bool:
-        """
-        Call the corresponding is_namespace function.
-
-        Parameters
-        ----------
-        xp : array_namespace
-            Array namespace to check.
-
-        Returns
-        -------
-        bool
-            ``True`` if xp matches the namespace, ``False`` otherwise.
-        """
-        is_namespace_func = getattr(_compat, f"is_{self.library_name}_namespace")
-        is_namespace_func = cast(Callable[[ModuleType], bool], is_namespace_func)
-        return is_namespace_func(xp)

src/array_api_extra/_lib/_utils/_compat.py

@@ -6,7 +6,9 @@
     from ...._array_api_compat_vendor import (  # pyright: ignore[reportMissingImports]
         array_namespace,
         device,
+        is_array_api_strict_namespace,
         is_cupy_namespace,
+        is_dask_namespace,
         is_jax_array,
         is_jax_namespace,
         is_numpy_namespace,
@@ -19,7 +21,9 @@
     from array_api_compat import (  # pyright: ignore[reportMissingTypeStubs]
         array_namespace,
         device,
+        is_array_api_strict_namespace,
         is_cupy_namespace,
+        is_dask_namespace,
         is_jax_array,
         is_jax_namespace,
         is_numpy_namespace,
@@ -32,7 +36,9 @@
 __all__ = [
     "array_namespace",
     "device",
+    "is_array_api_strict_namespace",
     "is_cupy_namespace",
+    "is_dask_namespace",
     "is_jax_array",
     "is_jax_namespace",
     "is_numpy_namespace",

src/array_api_extra/_lib/_utils/_compat.pyi

@@ -18,7 +18,9 @@ def array_namespace(
     use_compat: bool | None = None,
 ) -> ArrayModule: ...
 def device(x: Array, /) -> Device: ...
+def is_array_api_strict_namespace(xp: ModuleType, /) -> bool: ...
 def is_cupy_namespace(xp: ModuleType, /) -> bool: ...
+def is_dask_namespace(xp: ModuleType, /) -> bool: ...
 def is_jax_namespace(xp: ModuleType, /) -> bool: ...
 def is_numpy_namespace(xp: ModuleType, /) -> bool: ...
 def is_torch_namespace(xp: ModuleType, /) -> bool: ...

tests/test_funcs.py

@@ -421,7 +421,7 @@ def test_ndim(self, xp: ModuleType):
     def test_mode_not_implemented(self, xp: ModuleType):
         a = xp.arange(3)
         with pytest.raises(NotImplementedError, match="Only `'constant'`"):
-            pad(a, 2, mode="edge")  # type: ignore[arg-type]
+            pad(a, 2, mode="edge")  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
 
     def test_device(self, xp: ModuleType, device: Device):
         a = xp.asarray(0.0, device=device)