combine enums

Author: lucascolley

src/array_api_extra/_delegation.py

@@ -1,58 +1,15 @@
 """Delegation to existing implementations for Public API Functions."""
 
-import functools
-from enum import Enum
 from types import ModuleType
-from typing import final
 
-from ._lib import _funcs
-from ._lib._utils._compat import (
-    array_namespace,
-    is_cupy_namespace,
-    is_jax_namespace,
-    is_numpy_namespace,
-    is_torch_namespace,
-)
+from ._lib import Library, _funcs
+from ._lib._utils._compat import array_namespace
 from ._lib._utils._typing import Array
 
 __all__ = ["pad"]
 
 
-@final
-class IsNamespace(Enum):
-    """Enum to access is_namespace functions as the backend."""
-
-    # TODO: when Python 3.10 is dropped, use `enum.member`
-    # https://stackoverflow.com/a/74302109
-    CUPY = functools.partial(is_cupy_namespace)
-    JAX = functools.partial(is_jax_namespace)
-    NUMPY = functools.partial(is_numpy_namespace)
-    TORCH = functools.partial(is_torch_namespace)
-
-    def __call__(self, xp: ModuleType) -> bool:
-        """
-        Call the is_namespace function.
-
-        Parameters
-        ----------
-        xp : array_namespace
-            Array namespace to check.
-
-        Returns
-        -------
-        bool
-            ``True`` if xp matches the namespace, ``False`` otherwise.
-        """
-        return self.value(xp)
-
-
-CUPY = IsNamespace.CUPY
-JAX = IsNamespace.JAX
-NUMPY = IsNamespace.NUMPY
-TORCH = IsNamespace.TORCH
-
-
-def _delegate(xp: ModuleType, *backends: IsNamespace) -> bool:
+def _delegate(xp: ModuleType, *backends: Library) -> bool:
     """
     Check whether `xp` is one of the `backends` to delegate to.
 
@@ -68,7 +25,7 @@ def _delegate(xp: ModuleType, *backends: IsNamespace) -> bool:
     bool
         ``True`` if `xp` matches one of the `backends`, ``False`` otherwise.
     """
-    return any(is_namespace(xp) for is_namespace in backends)
+    return any(backend.is_namespace(xp) for backend in backends)
 
 
 def pad(
@@ -113,13 +70,13 @@ def pad(
         raise NotImplementedError(msg)
 
     # https://github.com/pytorch/pytorch/blob/cf76c05b4dc629ac989d1fb8e789d4fac04a095a/torch/_numpy/_funcs_impl.py#L2045-L2056
-    if _delegate(xp, TORCH):
+    if _delegate(xp, Library.TORCH):
         pad_width = xp.asarray(pad_width)
         pad_width = xp.broadcast_to(pad_width, (x.ndim, 2))
         pad_width = xp.flip(pad_width, axis=(0,)).flatten()
         return xp.nn.functional.pad(x, tuple(pad_width), value=constant_values)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
 
-    if _delegate(xp, NUMPY, JAX, CUPY):
+    if _delegate(xp, Library.NUMPY, Library.JAX_NUMPY, Library.CUPY):
         return xp.pad(x, pad_width, mode, constant_values=constant_values)
 
     return _funcs.pad(x, pad_width, constant_values=constant_values, xp=xp)

src/array_api_extra/_lib/__init__.py

@@ -1 +1,5 @@
 """Internals of array-api-extra."""
+
+from ._libraries import Library
+
+__all__ = ["Library"]

src/array_api_extra/_lib/_libraries.py

@@ -0,0 +1,69 @@
+"""Code specifying libraries array-api-extra interacts with."""
+
+from collections.abc import Callable
+from enum import Enum
+from types import ModuleType
+from typing import cast
+
+from ._utils import _compat
+
+__all__ = ["Library"]
+
+
+class Library(Enum):  # numpydoc ignore=PR01,PR02
+    """
+    All array library backends explicitly tested by array-api-extra.
+
+    Parameters
+    ----------
+    value : str
+        String describing the backend.
+    library_name : str
+        Name of the array library of 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"
+
+    def __new__(
+        cls, value: str, _library_name: str, _module_name: str
+    ):  # numpydoc ignore=GL08
+        obj = object.__new__(cls)
+        obj._value_ = value
+        return obj
+
+    def __init__(
+        self, _value: str, library_name: str, module_name: str
+    ):  # numpydoc ignore=GL08
+        self.library_name = library_name
+        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)

tests/conftest.py

@@ -1,33 +1,16 @@
 """Pytest fixtures."""
 
-from enum import Enum
 from types import ModuleType
 from typing import cast
 
 import pytest
 
+from array_api_extra._lib import Library
 from array_api_extra._lib._utils._compat import array_namespace
 from array_api_extra._lib._utils._compat import device as get_device
 from array_api_extra._lib._utils._typing import Device
 
 
-class Library(Enum):
-    """All array libraries explicitly tested by array-api-extra."""
-
-    ARRAY_API_STRICT = "array_api_strict"
-    NUMPY = "numpy"
-    NUMPY_READONLY = "numpy_readonly"
-    CUPY = "cupy"
-    TORCH = "torch"
-    DASK_ARRAY = "dask.array"
-    SPARSE = "sparse"
-    JAX_NUMPY = "jax.numpy"
-
-    def __str__(self) -> str:  # type: ignore[explicit-override]  # pyright: ignore[reportImplicitOverride]  # numpydoc ignore=RT01
-        """Pretty-print parameterized test names."""
-        return self.value
-
-
 @pytest.fixture(params=tuple(Library))
 def library(request: pytest.FixtureRequest) -> Library:  # numpydoc ignore=PR01,RT03
     """
@@ -60,8 +43,7 @@ def xp(library: Library) -> ModuleType:  # numpydoc ignore=PR01,RT03
     -------
     The current array namespace.
     """
-    name = "numpy" if library == Library.NUMPY_READONLY else library.value
-    xp = pytest.importorskip(name)
+    xp = pytest.importorskip(library.module_name)
     if library == Library.JAX_NUMPY:
         import jax  # type: ignore[import-not-found]  # pyright: ignore[reportMissingImports]
 

tests/test_at.py

@@ -12,12 +12,11 @@
 )
 
 from array_api_extra import at
+from array_api_extra._lib import Library
 from array_api_extra._lib._funcs import _AtOp
 from array_api_extra._lib._testing import xp_assert_equal
 from array_api_extra._lib._utils._typing import Array
 
-from .conftest import Library
-
 
 @pytest.fixture
 def array(library: Library, xp: ModuleType) -> Array:

tests/test_funcs.py

@@ -17,12 +17,11 @@
     setdiff1d,
     sinc,
 )
+from array_api_extra._lib import Library
 from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
 from array_api_extra._lib._utils._compat import device as get_device
 from array_api_extra._lib._utils._typing import Array, Device
 
-from .conftest import Library
-
 # mypy: disable-error-code=no-untyped-usage
 
 

tests/test_testing.py

@@ -1,10 +1,9 @@
 import numpy as np
 import pytest
 
+from array_api_extra._lib import Library
 from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
 
-from .conftest import Library
-
 # mypy: disable-error-code=no-any-decorated
 # pyright: reportUnknownParameterType=false,reportMissingParameterType=false
 

tests/test_utils.py

@@ -3,13 +3,12 @@
 import numpy as np
 import pytest
 
+from array_api_extra._lib import Library
 from array_api_extra._lib._testing import xp_assert_equal
 from array_api_extra._lib._utils._compat import device as get_device
 from array_api_extra._lib._utils._helpers import in1d
 from array_api_extra._lib._utils._typing import Array, Device
 
-from .conftest import Library
-
 # mypy: disable-error-code=no-untyped-usage