STY: reduce from imports some more

src/array_api_extra/_lib/_compat.py

@@ -1,4 +1,4 @@
-"""Acquire helpers from array-api-compat."""
+"""Helpers from array-api-compat."""
 # Allow packages that vendor both `array-api-extra` and
 # `array-api-compat` to override the import location
 

src/array_api_extra/_lib/_compat.pyi

@@ -1,6 +1,4 @@
-"""Static type stubs for `_compat.py`."""
-
-# https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
+"""Helpers from array-api-compat."""
 
 from typing import Any, TypeGuard
 

src/array_api_extra/_lib/_helpers.py

@@ -7,9 +7,9 @@
 import math
 import pickle
 import types
+import typing
 import warnings
 from collections.abc import Callable, Generator, Iterable, Iterator, Sequence
-from functools import wraps
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -19,18 +19,9 @@
     ParamSpec,
     TypeAlias,
     TypeVar,
-    cast,
 )
 
 from . import _compat
-from ._compat import (
-    array_namespace,
-    is_array_api_obj,
-    is_dask_namespace,
-    is_jax_namespace,
-    is_numpy_array,
-    is_torch_namespace,
-)
 from ._typing import Array, ArrayNamespace
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -102,7 +93,7 @@ def in1d(
     https://github.com/numpy/numpy/blob/v1.26.0/numpy/lib/arraysetops.py#L524-L758
     """
     if xp is None:
-        xp = array_namespace(x1, x2)
+        xp = _compat.array_namespace(x1, x2)
 
     x1_shape = eager_shape(x1)
     x2_shape = eager_shape(x2)
@@ -149,7 +140,7 @@ def is_python_scalar(x: object) -> TypeIs[complex]:  # numpydoc ignore=PR01,RT01
     # isinstance(x, float) returns True for np.float64
     # isinstance(x, complex) returns True for np.complex128
     # bool is a subclass of int
-    return isinstance(x, int | float | complex) and not is_numpy_array(x)
+    return isinstance(x, int | float | complex) and not _compat.is_numpy_array(x)
 
 
 def asarrays(
@@ -197,7 +188,7 @@ def asarrays(
         swap = True
         b, a = a, b
 
-    if is_array_api_obj(a):
+    if _compat.is_array_api_obj(a):
         # a is an Array API object
         # b is a int | float | complex | bool
         xa = a
@@ -209,7 +200,7 @@ def asarrays(
             float: ("real floating", "complex floating"),
             complex: "complex floating",
         }
-        kind = same_dtype[type(cast(complex, b))]
+        kind = same_dtype[type(typing.cast(complex, b))]
         if xp.isdtype(a.dtype, kind):
             xb = xp.asarray(b, dtype=a.dtype, device=_compat.device(a))
         else:
@@ -277,7 +268,7 @@ def eager_shape(x: Array, /, axis: int | None = None) -> tuple[int, ...]:
     if any(s is None or math.isnan(s) for s in shape):
         msg = "Unsupported lazy shape"
         raise TypeError(msg)
-    return cast(tuple[int, ...], shape)
+    return typing.cast(tuple[int, ...], shape)
 
 
 def meta_namespace(
@@ -301,12 +292,12 @@ def meta_namespace(
         If xp is Dask, the namespace of the Dask chunks;
         otherwise, the namespace of the arrays.
     """
-    xp = array_namespace(*arrays) if xp is None else xp
-    if not is_dask_namespace(xp):
+    xp = _compat.array_namespace(*arrays) if xp is None else xp
+    if not _compat.is_dask_namespace(xp):
         return xp
     # Quietly skip scalars and None's
-    metas = [cast(Array | None, getattr(a, "_meta", None)) for a in arrays]
-    return array_namespace(*metas)
+    metas = [typing.cast(Array | None, getattr(a, "_meta", None)) for a in arrays]
+    return _compat.array_namespace(*metas)
 
 
 def capabilities(xp: ArrayNamespace, x: Array | None = None) -> dict[str, int | None]:
@@ -330,13 +321,13 @@ def capabilities(xp: ArrayNamespace, x: Array | None = None) -> dict[str, int |
         Capabilities of the namespace.
     """
     out = xp.__array_namespace_info__().capabilities()
-    if is_jax_namespace(xp):
+    if _compat.is_jax_namespace(xp):
         if out["boolean indexing"]:  # pragma: no cover
             # Backwards compatibility with jax <0.6.0
             # https://github.com/jax-ml/jax/issues/27418
             out = out.copy()
             out["boolean indexing"] = False
-    elif is_torch_namespace(xp):
+    elif _compat.is_torch_namespace(xp):
         # FIXME https://github.com/data-apis/array-api/issues/945
         device = xp.get_default_device() if x is None else _compat.device(x)
         if device.type == "meta":  # type: ignore[union-attr]  # pyright: ignore[reportAttributeAccessIssue]
@@ -487,7 +478,7 @@ def pickle_unflatten(instances: Iterable[object], rest: FlattenRest) -> Any:
     or even the same types of objects. Excess elements, if any, will be left untouched.
     """
     iters = iter(instances), iter(rest)
-    pik = cast(bytes, next(iters[1]))
+    pik = typing.cast(bytes, next(iters[1]))
 
     class Unpickler(pickle.Unpickler):  # numpydoc ignore=GL08
         """Mirror of the overridden Pickler in pickle_flatten."""
@@ -602,7 +593,7 @@ def inner(  # numpydoc ignore=GL08
         res = func(*args, **kwargs)  # pyright: ignore[reportCallIssue]
         return _AutoJITWrapper(res)
 
-    @wraps(func)
+    @functools.wraps(func)
     def outer(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
         wargs = _AutoJITWrapper((args, kwargs))
         return inner(wargs).obj
@@ -622,8 +613,8 @@ def normalize_pad_width(
         and len(pad_width) == 2
         and all(isinstance(i, int) for i in pad_width)
     ):
-        return [cast(tuple[int, int], pad_width)] * ndim
-    return cast(list[tuple[int, int]], list(pad_width))
+        return [typing.cast(tuple[int, int], pad_width)] * ndim
+    return typing.cast(list[tuple[int, int]], list(pad_width))
 
 
 def is_jax_jit_enabled(xp: ArrayNamespace) -> bool:  # numpydoc ignore=PR01,RT01

src/array_api_extra/_lib/_testing.py

@@ -1,9 +1,9 @@
 """Private testing utilities."""
 
-from __future__ import annotations
-
 import pytest
 
+__all__ = ["xfail"]
+
 
 def xfail(
     request: pytest.FixtureRequest, *, reason: str, strict: bool | None = None

src/array_api_extra/_lib/_typing.py

@@ -1,6 +1,5 @@
 """Static typing helpers."""
 # numpydoc ignore=GL08
-# pylint: disable=duplicate-code
 
 from types import ModuleType