data-apis · lucascolley · Mar 20, 2025 · Mar 20, 2025
diff --git a/pixi.lock b/pixi.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -111,7 +111,7 @@ typing-extensions = "*"
 numpy = "*"
 
 [tool.pixi.feature.docs.tasks]
-docs = { cmd = "sphinx-build . build/", cwd = "docs" }
+docs = { cmd = "sphinx-build -E -W . build/", cwd = "docs" }
 open-docs = { cmd = "open build/index.html", cwd = "docs", depends-on = ["docs"] }
 
 [tool.pixi.feature.dev.dependencies]

diff --git a/src/array_api_extra/_lib/_at.py b/src/array_api_extra/_lib/_at.py
@@ -1,13 +1,12 @@
 """Update operations for read-only arrays."""
 
-# https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
 from __future__ import annotations
 
 import operator
 from collections.abc import Callable
 from enum import Enum
 from types import ModuleType
-from typing import ClassVar, cast
+from typing import TYPE_CHECKING, ClassVar, cast
 
 from ._utils._compat import (
     array_namespace,
@@ -18,6 +17,10 @@
 from ._utils._helpers import meta_namespace
 from ._utils._typing import Array, SetIndex
 
+if TYPE_CHECKING:  # pragma: no cover
+    # TODO import from typing (requires Python >=3.11)
+    from typing_extensions import Self
+
 
 class _AtOp(Enum):
     """Operations for use in `xpx.at`."""
@@ -204,7 +207,7 @@ def __init__(
         self._x = x
         self._idx = idx
 
-    def __getitem__(self, idx: SetIndex, /) -> at:  # numpydoc ignore=PR01,RT01
+    def __getitem__(self, idx: SetIndex, /) -> Self:  # numpydoc ignore=PR01,RT01
         """
         Allow for the alternate syntax ``at(x)[start:stop:step]``.
 
@@ -214,7 +217,7 @@ def __getitem__(self, idx: SetIndex, /) -> at:  # numpydoc ignore=PR01,RT01
         if self._idx is not _undef:
             msg = "Index has already been set"
             raise ValueError(msg)
-        return at(self._x, idx)
+        return type(self)(self._x, idx)
 
     def _op(
         self,

diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
@@ -1,8 +1,5 @@
 """Array-agnostic implementations for the public API."""
 
-# https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
-from __future__ import annotations
-
 import math
 import warnings
 from collections.abc import Callable, Sequence

diff --git a/src/array_api_extra/_lib/_lazy.py b/src/array_api_extra/_lib/_lazy.py
@@ -1,13 +1,12 @@
 """Public API Functions."""
 
-# https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
 from __future__ import annotations
 
 import math
 from collections.abc import Callable, Sequence
 from functools import partial, wraps
 from types import ModuleType
-from typing import TYPE_CHECKING, Any, cast, overload
+from typing import TYPE_CHECKING, Any, ParamSpec, TypeAlias, cast, overload
 
 from ._funcs import broadcast_shapes
 from ._utils import _compat
@@ -20,23 +19,15 @@
 from ._utils._typing import Array, DType
 
 if TYPE_CHECKING:  # pragma: no cover
-    # TODO move outside TYPE_CHECKING
-    # depends on scikit-learn abandoning Python 3.9
-    # https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
-    from typing import ParamSpec, TypeAlias
-
     import numpy as np
     from numpy.typing import ArrayLike
 
     NumPyObject: TypeAlias = np.ndarray[Any, Any] | np.generic  # type: ignore[explicit-any]
-    P = ParamSpec("P")
 else:
-    # Sphinx hacks
+    # Sphinx hack
     NumPyObject = Any
 
-    class P:  # pylint: disable=missing-class-docstring
-        args: tuple
-        kwargs: dict
+P = ParamSpec("P")
 
 
 @overload

diff --git a/src/array_api_extra/_lib/_utils/_helpers.py b/src/array_api_extra/_lib/_utils/_helpers.py
@@ -1,6 +1,5 @@
 """Helper functions used by `array_api_extra/_funcs.py`."""
 
-# https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
 from __future__ import annotations
 
 import math

diff --git a/src/array_api_extra/testing.py b/src/array_api_extra/testing.py
@@ -4,42 +4,33 @@
 See also _lib._testing for additional private testing utilities.
 """
 
-# https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
 from __future__ import annotations
 
 import contextlib
 from collections.abc import Callable, Iterable, Iterator, Sequence
 from functools import wraps
 from types import ModuleType
-from typing import TYPE_CHECKING, Any, TypeVar, cast
+from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast
 
 from ._lib._utils._compat import is_dask_namespace, is_jax_namespace
 
 __all__ = ["lazy_xp_function", "patch_lazy_xp_functions"]
 
 if TYPE_CHECKING:  # pragma: no cover
-    # TODO move ParamSpec outside TYPE_CHECKING
-    # depends on scikit-learn abandoning Python 3.9
-    # https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
-    from typing import ParamSpec
-
+    # TODO import override from typing (requires Python >=3.12)
     import pytest
     from dask.typing import Graph, Key, SchedulerGetCallable
     from typing_extensions import override
 
-    P = ParamSpec("P")
 else:
-    SchedulerGetCallable = object
-
     # Sphinx hacks
-    class P:  # pylint: disable=missing-class-docstring
-        args: tuple
-        kwargs: dict
+    SchedulerGetCallable = object
 
-    def override(func: Callable[P, T]) -> Callable[P, T]:
+    def override(func: object) -> object:
         return func
 
 
+P = ParamSpec("P")
 T = TypeVar("T")
 
 _ufuncs_tags: dict[object, dict[str, Any]] = {}  # type: ignore[explicit-any]