repr: use has_xp() for array-api detection in ArrayAPIFormatter

Two-tier detection: tier 1 uses the canonical has_xp() protocol check
from anndata.compat (catches JAX, numpy >=2.0); tier 2 falls back to
duck-typing (shape/dtype/ndim) for arrays that don't yet implement the
full protocol (PyTorch, TensorFlow). Also uses __array_namespace__()
for backend label resolution and updates stale PR scverse#2063 → scverse#2071.

Author: katosh

src/anndata/_repr/formatters.py

@@ -45,6 +45,7 @@
     CSS_NESTED_ANNDATA,
     CSS_TEXT_MUTED,
 )
+from ..compat import has_xp
 from .components import render_category_list
 from .lazy import get_lazy_categorical_info, is_lazy_column
 from .registry import (
@@ -72,6 +73,8 @@
 )
 
 if TYPE_CHECKING:
+    from typing import ClassVar
+
     from .registry import FormatterContext
 
 
@@ -765,75 +768,77 @@ class ArrayAPIFormatter(TypeFormatter):
     """
     Formatter for Array-API compatible arrays (JAX, PyTorch, TensorFlow, etc.).
 
-    Future-proofing notes:
-    - PR #2063 (https://github.com/scverse/anndata/pull/2063) adds Array-API compatibility
-    - Handles JAX arrays, PyTorch tensors, TensorFlow tensors, and other array-like objects
-    - Uses duck typing to detect array-like objects without isinstance checks
-    - Low priority (50) ensures specific formatters (numpy, cupy, etc.) are tried first
+    Detection strategy (two tiers):
 
-    References:
-    - Array API Standard: https://data-apis.org/array-api/latest/
+    1. :func:`~anndata.compat.has_xp` — canonical check for arrays implementing the
+       `Array API standard <https://data-apis.org/array-api/latest/>`_ (e.g. JAX, CuPy).
+    2. Duck-typing fallback — catches arrays with ``shape``/``dtype``/``ndim`` that do
+       not (yet) implement the full protocol (e.g. PyTorch tensors, TensorFlow tensors).
+
+    Low priority (50) ensures specific formatters (numpy, cupy, dask, etc.)
+    are tried first.
     """
 
     priority = 50  # Lower than specific formatters (numpy=110, cupy=120) but higher than builtins
 
+    _FRIENDLY_NAMES: ClassVar[dict[str, str]] = {
+        "jax": "JAX",
+        "jaxlib": "JAX",
+        "torch": "PyTorch",
+        "tensorflow": "TensorFlow",
+        "tf": "TensorFlow",
+        "mxnet": "MXNet",
+        "cupy": "CuPy",
+    }
+
+    # Modules already handled by dedicated formatters (defensive guard;
+    # the priority system normally prevents reaching this formatter).
+    _HANDLED_MODULES: ClassVar[tuple[str, ...]] = (
+        "numpy",
+        "pandas",
+        "scipy.sparse",
+        "cupy",
+        "cupyx",
+        "awkward",
+        "dask",
+    )
+
     def can_format(self, obj: object, context: FormatterContext) -> bool:
-        # Duck typing: Check for array-like attributes
-        # Must have shape, dtype, and ndim (array-api standard)
-        has_array_attrs = (
-            hasattr(obj, "shape") and hasattr(obj, "dtype") and hasattr(obj, "ndim")
-        )
+        # Tier 1: full Array API protocol (JAX, CuPy ≥12, numpy ≥2.0, …)
+        if has_xp(obj):
+            # numpy has its own formatter
+            return not isinstance(obj, np.ndarray)
 
-        if not has_array_attrs:
+        # Tier 2: duck-typing for arrays that expose shape/dtype/ndim
+        # but don't implement the full protocol (PyTorch, TensorFlow, …)
+        if not (
+            hasattr(obj, "shape") and hasattr(obj, "dtype") and hasattr(obj, "ndim")
+        ):
             return False
 
-        # Exclude types that already have specific formatters
-        # This prevents conflicts and ensures more specific formatters take precedence
+        # Exclude types that have dedicated formatters
         module = type(obj).__module__
-        already_handled = isinstance(
-            obj, (np.ndarray, pd.DataFrame, pd.Series)
-        ) or module.startswith((
-            "numpy",
-            "pandas",
-            "scipy.sparse",
-            "cupy",  # Has CuPyArrayFormatter
-            "cupyx",
-            "awkward",  # Has AwkwardArrayFormatter
-            "dask",  # Has DaskArrayFormatter
-        ))
-
-        return not already_handled
+        return not isinstance(obj, np.ndarray) and not module.startswith(
+            self._HANDLED_MODULES
+        )
 
     def format(self, obj: object, context: FormatterContext) -> FormattedOutput:
-        # Extract module and type information
-        module_name = type(obj).__module__.split(".")[0]  # e.g., "jax" from "jax.numpy"
         type_name = type(obj).__name__
-
         shape_str = " × ".join(format_number(s) for s in obj.shape)
         dtype_str = str(obj.dtype)
 
-        # Detect common array-api backends
-        # Known backends: JAX, PyTorch, TensorFlow, MXNet, etc.
-        known_backends = {
-            "jax": "JAX",
-            "jaxlib": "JAX",
-            "torch": "PyTorch",
-            "tensorflow": "TensorFlow",
-            "tf": "TensorFlow",
-            "mxnet": "MXNet",
-        }
-        backend_label = known_backends.get(module_name, module_name)
-
-        # Try to get device information (GPU vs CPU)
+        # Derive backend label: prefer __array_namespace__, fall back to module name
+        backend_label = type(obj).__module__.split(".")[0]
+        with contextlib.suppress(Exception):
+            xp = obj.__array_namespace__()
+            ns_name = getattr(xp, "__name__", "") or type(xp).__module__
+            backend_label = ns_name.split(".")[0]
+        backend_label = self._FRIENDLY_NAMES.get(backend_label, backend_label)
+
+        # Device info (present on array-api arrays; also on PyTorch/CuPy)
         device_info = ""
-        try:
-            if hasattr(obj, "device"):
-                device_info = f" on {obj.device}"
-            elif hasattr(obj, "device_buffer"):  # JAX
-                device_info = f" on {obj.device_buffer.device()}"
-        except Exception:  # noqa: BLE001
-            # Intentional broad catch: device access varies by backend and can fail
-            pass
+        with contextlib.suppress(Exception):
+            device_info = f" on {obj.device}"
 
         # For obsm/varm sections, show number of columns in preview
         preview = None
@@ -1126,7 +1131,7 @@ def _register_builtin_formatters() -> None:
         LazyColumnFormatter(),
         # Low-medium priority (Array-API compatible arrays)
         # Must come after specific array formatters (numpy, cupy, etc.) but before builtins
-        # Handles JAX, PyTorch, TensorFlow arrays added in PR #2063
+        # Handles JAX, PyTorch, TensorFlow arrays (PR #2071 added array-api support)
         ArrayAPIFormatter(),
         # Low priority (builtins)
         NoneFormatter(),

tests/repr/test_repr_formatters.py

@@ -20,6 +20,33 @@
 from .conftest import HAS_AWKWARD, HAS_DASK
 
 
+def _make_array_api_mock(module: str, *, shape, dtype, device="cpu"):
+    """Create a mock array satisfying the SupportsArrayApi protocol.
+
+    The mock has all attributes required by ``anndata.types.SupportsArrayApi``
+    (``shape``, ``device``, ``__array_namespace__``, ``to_device``,
+    ``__dlpack__``, ``__dlpack_device__``) so that ``has_xp()`` returns True.
+    """
+    ns_module = type("Namespace", (), {"__name__": module.split(".")[0]})()
+
+    cls = type(
+        "MockArrayAPI",
+        (),
+        {
+            "shape": shape,
+            "dtype": dtype,
+            "ndim": len(shape),
+            "device": device,
+            "__array_namespace__": lambda self, **kw: ns_module,
+            "to_device": lambda self, dev, /, **kw: self,
+            "__dlpack__": lambda self, **kw: None,
+            "__dlpack_device__": lambda self: (1, 0),
+        },
+    )
+    cls.__module__ = module
+    return cls()
+
+
 class TestNumpyFormatters:
     """Tests for NumPy array formatters."""
 
@@ -506,39 +533,35 @@ def test_array_api_formatter_jax_like(self):
         from anndata._repr.formatters import ArrayAPIFormatter
         from anndata._repr.registry import FormatterContext
 
-        class MockJAXArray:
-            def __init__(self):
-                self.shape = (100, 50)
-                self.dtype = np.float32
-                self.ndim = 2
-                self.device = "gpu:0"
-
-        MockJAXArray.__module__ = "jax.numpy"
+        mock_arr = _make_array_api_mock(
+            "jax.numpy", shape=(100, 50), dtype=np.float32, device="gpu:0"
+        )
 
         formatter = ArrayAPIFormatter()
-        mock_arr = MockJAXArray()
 
         assert formatter.can_format(mock_arr, FormatterContext())
         result = formatter.format(mock_arr, FormatterContext())
 
-        assert "MockJAXArray" in result.type_name
         assert "100" in result.type_name
         assert "50" in result.type_name
         assert result.css_class == "anndata-dtype--array-api"
-        assert "JAX" in result.tooltip
+        assert "JAX" in result.tooltip or "jax" in result.tooltip.lower()
         assert "gpu:0" in result.tooltip
 
-    def test_array_api_formatter_pytorch_like(self):
-        """Test Array-API formatter with PyTorch-like tensor."""
+    def test_array_api_formatter_pytorch_like_duck_typing(self):
+        """Test Array-API formatter with PyTorch-like tensor (duck-typing tier).
+
+        PyTorch doesn't implement __array_namespace__, so it falls back
+        to the duck-typing check (shape/dtype/ndim).
+        """
         from anndata._repr.formatters import ArrayAPIFormatter
         from anndata._repr.registry import FormatterContext
 
         class MockTorchTensor:
-            def __init__(self):
-                self.shape = (64, 32)
-                self.dtype = "torch.float32"
-                self.ndim = 2
-                self.device = "cuda:0"
+            shape = (64, 32)
+            dtype = "torch.float32"
+            ndim = 2
+            device = "cuda:0"
 
         MockTorchTensor.__module__ = "torch"
 
@@ -548,34 +571,24 @@ def __init__(self):
         assert formatter.can_format(mock_tensor, FormatterContext())
         result = formatter.format(mock_tensor, FormatterContext())
 
-        assert "MockTorchTensor" in result.type_name
-        assert "PyTorch" in result.tooltip
+        assert "PyTorch" in result.tooltip or "torch" in result.tooltip.lower()
+        assert "cuda:0" in result.tooltip
 
-    def test_array_api_formatter_device_buffer(self):
-        """Test Array-API formatter with device_buffer attribute."""
+    def test_array_api_formatter_protocol_array(self):
+        """Test Array-API formatter with full protocol array (tier 1)."""
         from anndata._repr.formatters import ArrayAPIFormatter
         from anndata._repr.registry import FormatterContext
 
-        class MockDeviceBuffer:
-            def device(self):
-                return "tpu:0"
-
-        class MockJAXArrayWithBuffer:
-            def __init__(self):
-                self.shape = (10, 5)
-                self.dtype = np.float32
-                self.ndim = 2
-                self.device_buffer = MockDeviceBuffer()
-
-        MockJAXArrayWithBuffer.__module__ = "jaxlib.xla_extension"
+        mock_arr = _make_array_api_mock(
+            "jax.numpy", shape=(10, 5), dtype=np.float32, device="tpu:0"
+        )
 
         formatter = ArrayAPIFormatter()
-        mock_arr = MockJAXArrayWithBuffer()
 
         assert formatter.can_format(mock_arr, FormatterContext())
         result = formatter.format(mock_arr, FormatterContext())
 
-        assert "JAX" in result.tooltip
+        assert "JAX" in result.tooltip or "jax" in result.tooltip.lower()
         assert "tpu:0" in result.tooltip
 
     def test_array_api_formatter_excludes_numpy(self):
@@ -677,22 +690,9 @@ def test_array_api_formatter_with_mock_jax_array(self):
         from anndata._repr.formatters import ArrayAPIFormatter
         from anndata._repr.registry import FormatterContext
 
-        class MockJAXArray:
-            def __init__(self):
-                self.shape = (100, 50)
-                self.dtype = np.dtype("float32")
-                self.ndim = 2
-
-            @property
-            def __module__(self):
-                return "jax.numpy"
-
-            def __class__(self):
-                return type("DeviceArray", (), {})
-
-        mock_array = MockJAXArray()
-        type(mock_array).__module__ = "jax.numpy"
-        type(mock_array).__name__ = "DeviceArray"
+        mock_array = _make_array_api_mock(
+            "jax.numpy", shape=(100, 50), dtype=np.dtype("float32"), device="cpu"
+        )
 
         formatter = ArrayAPIFormatter()
         can_format = formatter.can_format(mock_array, FormatterContext())
@@ -987,17 +987,11 @@ def test_array_api_formatter_obsm_preview(self):
         from anndata._repr.formatters import ArrayAPIFormatter
         from anndata._repr.registry import FormatterContext
 
-        # Mock JAX-like array
-        class MockJAXArray:
-            def __init__(self):
-                self.shape = (100, 15)
-                self.dtype = np.float32
-                self.ndim = 2
-
-        MockJAXArray.__module__ = "jax.numpy"
+        arr = _make_array_api_mock(
+            "jax.numpy", shape=(100, 15), dtype=np.float32, device="cpu"
+        )
 
         formatter = ArrayAPIFormatter()
-        arr = MockJAXArray()
 
         # No preview outside obsm/varm
         result = formatter.format(arr, FormatterContext(section="uns"))