From 89992149be16778249da5defa5054ad9273e425b Mon Sep 17 00:00:00 2001
From: Lucas Colley <lucas.colley8@gmail.com>
Date: Sat, 15 Aug 2026 13:46:54 +0100
Subject: [PATCH] MAINT: hide some typing hacks away

---
 src/array_api_extra/_lib/_helpers.py    | 12 +---------
 src/array_api_extra/_lib/_typing.py     | 20 ++++++++++++++--
 src/array_api_extra/_lib/_typing.pyi    | 10 ++++++--
 src/array_api_extra/testing/_testing.py | 32 ++++++++++++-------------
 tests/test_helpers.py                   | 13 ++--------
 5 files changed, 44 insertions(+), 43 deletions(-)

diff --git a/src/array_api_extra/_lib/_helpers.py b/src/array_api_extra/_lib/_helpers.py
index 60669901..9cf74b53 100644
--- a/src/array_api_extra/_lib/_helpers.py
+++ b/src/array_api_extra/_lib/_helpers.py
@@ -11,7 +11,6 @@
 import warnings
 from collections.abc import Callable, Generator, Iterable, Iterator, Sequence
 from typing import (
-    TYPE_CHECKING,
     Any,
     ClassVar,
     Generic,
@@ -22,16 +21,7 @@
 )
 
 from . import _compat
-from ._typing import Array, ArrayNamespace
-
-if TYPE_CHECKING:  # pragma: no cover
-    # TODO import from typing (requires Python >=3.12 and >=3.13)
-    from typing_extensions import TypeIs, override
-else:
-
-    def override(func):
-        return func
-
+from ._typing import Array, ArrayNamespace, TypeIs, override
 
 P = ParamSpec("P")
 T = TypeVar("T")
diff --git a/src/array_api_extra/_lib/_typing.py b/src/array_api_extra/_lib/_typing.py
index cd5e0c7e..240add20 100644
--- a/src/array_api_extra/_lib/_typing.py
+++ b/src/array_api_extra/_lib/_typing.py
@@ -1,23 +1,39 @@
-"""Static typing helpers."""
-# numpydoc ignore=GL08
+"""Static typing helpers — better implementations in the stub file."""
 
+from collections.abc import Callable
 from types import ModuleType
+from typing import Any
 
 Array = object
 ArrayLike = object
 ArrayNamespace = ModuleType
 DType = object
 Device = object
+Key = object
 GetIndex = object
+Graph = object
 NumPyObject = object
+SchedulerGetCallable = object
 SetIndex = object
 
+TypeIs = Any
+
 __all__ = [
     "Array",
+    "ArrayLike",
     "ArrayNamespace",
     "DType",
     "Device",
     "GetIndex",
+    "Graph",
+    "Key",
     "NumPyObject",
+    "SchedulerGetCallable",
     "SetIndex",
+    "TypeIs",
+    "override",
 ]
+
+
+def override(func: Callable[..., Any]) -> Callable[..., Any]:  # numpydoc ignore=GL08
+    return func
diff --git a/src/array_api_extra/_lib/_typing.pyi b/src/array_api_extra/_lib/_typing.pyi
index 67137946..0cab9da9 100644
--- a/src/array_api_extra/_lib/_typing.pyi
+++ b/src/array_api_extra/_lib/_typing.pyi
@@ -4,10 +4,11 @@ from types import EllipsisType, ModuleType
 from typing import Any, Protocol, TypeAlias
 
 import numpy as np
+from dask.typing import Graph, Key, SchedulerGetCallable
 from numpy.typing import ArrayLike
 
-# TODO import from typing (requires Python >=3.12)
-from typing_extensions import override
+# TODO import from typing (requires Python >=3.12 and >=3.13)
+from typing_extensions import TypeIs, override
 
 # TODO: use array-api-typing once it is available
 
@@ -112,6 +113,11 @@ __all__ = [
     "DType",
     "Device",
     "GetIndex",
+    "Graph",
+    "Key",
     "NumPyObject",
+    "SchedulerGetCallable",
     "SetIndex",
+    "TypeIs",
+    "override",
 ]
diff --git a/src/array_api_extra/testing/_testing.py b/src/array_api_extra/testing/_testing.py
index 2e30c1c9..ff1f1cc6 100644
--- a/src/array_api_extra/testing/_testing.py
+++ b/src/array_api_extra/testing/_testing.py
@@ -11,10 +11,22 @@
 import warnings
 from collections.abc import Callable, Generator, Iterator, Sequence
 from types import FunctionType, ModuleType
-from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar
+from typing import Any, ParamSpec, TypeVar
 
 from .._lib import _compat, _helpers
-from .._lib._typing import Array, ArrayNamespace, Device
+from .._lib._typing import (
+    Array,
+    ArrayNamespace,
+    Device,
+    Graph,
+    Key,
+    SchedulerGetCallable,
+    override,
+)
+
+if typing.TYPE_CHECKING:
+    import numpy as np
+    import pytest
 
 __all__ = [
     "assert_close",
@@ -25,20 +37,6 @@
     "patch_lazy_xp_functions",
 ]
 
-if TYPE_CHECKING:  # pragma: no cover
-    # TODO import override from typing (requires Python >=3.12)
-    import numpy as np
-    import pytest
-    from dask.typing import Graph, Key, SchedulerGetCallable
-    from typing_extensions import override as _override
-
-else:
-    # Sphinx hacks
-    SchedulerGetCallable = object
-
-    def _override(func):
-        return func
-
 
 __all__ = [
     "assert_close",
@@ -493,7 +491,7 @@ def __init__(self, max_count: int, msg: str) -> None:  # numpydoc ignore=GL08
         self.max_count = max_count
         self.msg = msg
 
-    @_override
+    @override
     def __call__(
         self, dsk: Graph, keys: Sequence[Key] | Key, **kwargs: Any
     ) -> Any:  # numpydoc ignore=GL08
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
index a9904456..34e1bc8c 100644
--- a/tests/test_helpers.py
+++ b/tests/test_helpers.py
@@ -1,5 +1,5 @@
 from collections.abc import Iterator
-from typing import TYPE_CHECKING, Generic, TypeVar, cast
+from typing import Generic, TypeVar, cast
 
 import numpy as np
 import pytest
@@ -18,20 +18,11 @@
     pickle_flatten,
     pickle_unflatten,
 )
-from array_api_extra._lib._typing import Array, ArrayNamespace, Device, DType
+from array_api_extra._lib._typing import Array, ArrayNamespace, Device, DType, override
 from array_api_extra.testing import assert_equal, lazy_xp_function
 
 from .conftest import np_compat
 
-if TYPE_CHECKING:  # pragma: no cover
-    # TODO import from typing (requires Python >=3.12)
-    from typing_extensions import override
-else:
-
-    def override(func):
-        return func
-
-
 T = TypeVar("T")
 
 # FIXME calls xp.unique_values without size