diff --git a/src/numpy-stubs/lib/_ufunclike_impl.pyi b/src/numpy-stubs/lib/_ufunclike_impl.pyi index 46ddc00b..0d1d3499 100644 --- a/src/numpy-stubs/lib/_ufunclike_impl.pyi +++ b/src/numpy-stubs/lib/_ufunclike_impl.pyi @@ -1,30 +1,100 @@ -from typing import Any, overload +from collections.abc import Sequence +from typing import overload from typing_extensions import TypeVar import numpy as np -from numpy._typing import NDArray, _ArrayLikeFloat_co, _ArrayLikeObject_co, _FloatLike_co +from _numtype import ( + Array, + Array_1d, + Array_2d, + Array_3d, + CanArraySized, + CoFloating_0d, + CoFloating_1ds, + CoFloating_1nd, + CoFloating_2ds, + CoFloating_3ds, + CoFloating_nd, + Is, + ToBool_0d, + ToBool_1nd, + ToFloat64_0d, + ToFloat64_1nd, + ToIntP_0d, + ToIntP_1nd, + _ToArray1_0d, + _ToArray1_1ds, + _ToArray1_1nd, + _ToArray1_2ds, + _ToArray1_3ds, +) __all__ = ["fix", "isneginf", "isposinf"] -_ArrayType = TypeVar("_ArrayType", bound=NDArray[Any]) +### +_ArrayT = TypeVar("_ArrayT", bound=Array) +_ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...]) +_RealT = TypeVar("_RealT", bound=np.integer | np.floating | np.object_) + +### + +# +@overload +def fix(x: CoFloating_nd, out: _ArrayT) -> _ArrayT: ... +@overload +def fix(x: _ArrayT, out: None) -> _ArrayT: ... +@overload +def fix(x: ToBool_0d, out: None = None) -> np.bool: ... # type: ignore[overload-overlap] +@overload +def fix(x: ToIntP_0d, out: None = None) -> np.intp: ... # type: ignore[overload-overlap] +@overload +def fix(x: ToFloat64_0d, out: None = None) -> np.float64: ... # type: ignore[overload-overlap] +@overload +def fix(x: _ToArray1_0d[_RealT], out: None = None) -> _RealT: ... # type: ignore[overload-overlap] +@overload +def fix(x: _ToArray1_1ds[_RealT], out: None = None) -> Array_1d[_RealT]: ... +@overload +def fix(x: _ToArray1_2ds[_RealT], out: None = None) -> Array_2d[_RealT]: ... @overload -def fix(x: _FloatLike_co, out: None = ...) -> np.floating: ... +def fix(x: _ToArray1_3ds[_RealT], out: None = None) -> Array_3d[_RealT]: ... @overload -def fix(x: _ArrayLikeFloat_co, out: None = ...) -> NDArray[np.floating]: ... +def fix(x: ToBool_1nd, out: None = None) -> Array[np.bool]: ... @overload -def fix(x: _ArrayLikeObject_co, out: None = ...) -> NDArray[np.object_]: ... +def fix(x: ToIntP_1nd, out: None = None) -> Array[np.intp]: ... +@overload # the `Sequence` is a workaround for a pyright promotion bug +def fix(x: ToFloat64_1nd | Sequence[Is[float]], out: None = None) -> Array[np.float64]: ... +@overload +def fix(x: _ToArray1_1nd[_RealT], out: None = None) -> Array[_RealT]: ... + +# +@overload +def isposinf(x: CoFloating_0d, out: None = None) -> np.bool: ... # type: ignore[overload-overlap] +@overload +def isposinf(x: CanArraySized[np.bool | np.integer | np.floating, _ShapeT], out: None = None) -> Array[np.bool, _ShapeT]: ... # type: ignore[overload-overlap] +@overload +def isposinf(x: CoFloating_1ds, out: None = None) -> Array_1d[np.bool]: ... +@overload +def isposinf(x: CoFloating_2ds, out: None = None) -> Array_2d[np.bool]: ... +@overload +def isposinf(x: CoFloating_3ds, out: None = None) -> Array_3d[np.bool]: ... +@overload +def isposinf(x: CoFloating_1nd, out: None = None) -> Array[np.bool]: ... +@overload +def isposinf(x: CoFloating_nd, out: _ArrayT) -> _ArrayT: ... + +# @overload -def fix(x: _ArrayLikeFloat_co | _ArrayLikeObject_co, out: _ArrayType) -> _ArrayType: ... +def isneginf(x: CoFloating_0d, out: None = None) -> np.bool: ... # type: ignore[overload-overlap] @overload -def isposinf(x: _FloatLike_co, out: None = ...) -> np.bool: ... +def isneginf(x: CanArraySized[np.bool | np.integer | np.floating, _ShapeT], out: None = None) -> Array[np.bool, _ShapeT]: ... # type: ignore[overload-overlap] @overload -def isposinf(x: _ArrayLikeFloat_co, out: None = ...) -> NDArray[np.bool]: ... +def isneginf(x: CoFloating_1ds, out: None = None) -> Array_1d[np.bool]: ... @overload -def isposinf(x: _ArrayLikeFloat_co, out: _ArrayType) -> _ArrayType: ... +def isneginf(x: CoFloating_2ds, out: None = None) -> Array_2d[np.bool]: ... @overload -def isneginf(x: _FloatLike_co, out: None = ...) -> np.bool: ... +def isneginf(x: CoFloating_3ds, out: None = None) -> Array_3d[np.bool]: ... @overload -def isneginf(x: _ArrayLikeFloat_co, out: None = ...) -> NDArray[np.bool]: ... +def isneginf(x: CoFloating_1nd, out: None = None) -> Array[np.bool]: ... @overload -def isneginf(x: _ArrayLikeFloat_co, out: _ArrayType) -> _ArrayType: ... +def isneginf(x: CoFloating_nd, out: _ArrayT) -> _ArrayT: ... diff --git a/test/static/accept/ufunclike.pyi b/test/static/accept/ufunclike.pyi index fdee886a..32ff5988 100644 --- a/test/static/accept/ufunclike.pyi +++ b/test/static/accept/ufunclike.pyi @@ -1,31 +1,36 @@ -from typing_extensions import assert_type +from typing import TypeAlias +from typing_extensions import TypeVar, assert_type import numpy as np import numpy.typing as npt +AR_u: npt.NDArray[np.uint32] +AR_U: npt.NDArray[np.str_] AR_LIKE_b: list[bool] -AR_LIKE_u: list[np.uint32] AR_LIKE_i: list[int] AR_LIKE_f: list[float] -AR_LIKE_O: list[np.object_] -AR_U: npt.NDArray[np.str_] +### + +_ScalarT = TypeVar("_ScalarT", bound=np.generic) +_Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_ScalarT]] + +### -assert_type(np.fix(AR_LIKE_b), npt.NDArray[np.floating]) -assert_type(np.fix(AR_LIKE_u), npt.NDArray[np.floating]) -assert_type(np.fix(AR_LIKE_i), npt.NDArray[np.floating]) -assert_type(np.fix(AR_LIKE_f), npt.NDArray[np.floating]) -assert_type(np.fix(AR_LIKE_O), npt.NDArray[np.object_]) +assert_type(np.fix(AR_u), npt.NDArray[np.uint32]) +assert_type(np.fix(AR_LIKE_b), npt.NDArray[np.bool]) +assert_type(np.fix(AR_LIKE_i), npt.NDArray[np.intp]) +assert_type(np.fix(AR_LIKE_f), npt.NDArray[np.float64]) assert_type(np.fix(AR_LIKE_f, out=AR_U), npt.NDArray[np.str_]) -assert_type(np.isposinf(AR_LIKE_b), npt.NDArray[np.bool]) -assert_type(np.isposinf(AR_LIKE_u), npt.NDArray[np.bool]) -assert_type(np.isposinf(AR_LIKE_i), npt.NDArray[np.bool]) -assert_type(np.isposinf(AR_LIKE_f), npt.NDArray[np.bool]) +assert_type(np.isposinf(AR_u), npt.NDArray[np.bool]) +assert_type(np.isposinf(AR_LIKE_b), _Array1D[np.bool]) +assert_type(np.isposinf(AR_LIKE_i), _Array1D[np.bool]) +assert_type(np.isposinf(AR_LIKE_f), _Array1D[np.bool]) assert_type(np.isposinf(AR_LIKE_f, out=AR_U), npt.NDArray[np.str_]) -assert_type(np.isneginf(AR_LIKE_b), npt.NDArray[np.bool]) -assert_type(np.isneginf(AR_LIKE_u), npt.NDArray[np.bool]) -assert_type(np.isneginf(AR_LIKE_i), npt.NDArray[np.bool]) -assert_type(np.isneginf(AR_LIKE_f), npt.NDArray[np.bool]) +assert_type(np.isneginf(AR_u), npt.NDArray[np.bool]) +assert_type(np.isneginf(AR_LIKE_b), _Array1D[np.bool]) +assert_type(np.isneginf(AR_LIKE_i), _Array1D[np.bool]) +assert_type(np.isneginf(AR_LIKE_f), _Array1D[np.bool]) assert_type(np.isneginf(AR_LIKE_f, out=AR_U), npt.NDArray[np.str_])