diff --git a/src/array_api_extra/_lib/_compat.py b/src/array_api_extra/_lib/_compat.py index 98f213ba..17300805 100644 --- a/src/array_api_extra/_lib/_compat.py +++ b/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 diff --git a/src/array_api_extra/_lib/_compat.pyi b/src/array_api_extra/_lib/_compat.pyi index cc9a6c33..a3b5f317 100644 --- a/src/array_api_extra/_lib/_compat.pyi +++ b/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 diff --git a/src/array_api_extra/_lib/_helpers.py b/src/array_api_extra/_lib/_helpers.py index af69c378..60669901 100644 --- a/src/array_api_extra/_lib/_helpers.py +++ b/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 diff --git a/src/array_api_extra/_lib/_testing.py b/src/array_api_extra/_lib/_testing.py index 9dbc3943..baa451f2 100644 --- a/src/array_api_extra/_lib/_testing.py +++ b/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 diff --git a/src/array_api_extra/_lib/_typing.py b/src/array_api_extra/_lib/_typing.py index 56efb87f..cd5e0c7e 100644 --- a/src/array_api_extra/_lib/_typing.py +++ b/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