Merge pull request #20199 from BvB93/numeric

ENH: Add dtype typing support to `np.core.numeric`

1 files changed, 473 insertions, 59 deletions

diff --git a/numpy/core/numeric.pyi b/numpy/core/numeric.pyi
index 54ab4b7c8..d7ec30351 100644
--- a/numpy/core/numeric.pyi
+++ b/numpy/core/numeric.pyi
@@ -1,6 +1,5 @@
 from typing import (
     Any,
-    Optional,
     Union,
     Sequence,
     Tuple,
@@ -8,18 +7,64 @@ from typing import (
     List,
     overload,
     TypeVar,
-    Iterable,
     Literal,
+    Type,
+    SupportsAbs,
+    SupportsIndex,
+    NoReturn,
 )
+from typing_extensions import TypeGuard
 
-from numpy import ndarray, generic, dtype, bool_, signedinteger, _OrderKACF, _OrderCF
-from numpy.typing import ArrayLike, DTypeLike, _ShapeLike
+from numpy import (
+    ComplexWarning as ComplexWarning,
+    dtype,
+    generic,
+    unsignedinteger,
+    signedinteger,
+    floating,
+    complexfloating,
+    bool_,
+    int_,
+    intp,
+    float64,
+    timedelta64,
+    object_,
+    _OrderKACF,
+    _OrderCF,
+)
+
+from numpy.typing import (
+    ArrayLike,
+    NDArray,
+    DTypeLike,
+    _ShapeLike,
+    _SupportsDType,
+    _FiniteNestedSequence,
+    _SupportsArray,
+    _ScalarLike_co,
+    _ArrayLikeBool_co,
+    _ArrayLikeUInt_co,
+    _ArrayLikeInt_co,
+    _ArrayLikeFloat_co,
+    _ArrayLikeComplex_co,
+    _ArrayLikeTD64_co,
+    _ArrayLikeObject_co,
+)
 
 _T = TypeVar("_T")
-_ArrayType = TypeVar("_ArrayType", bound=ndarray)
+_SCT = TypeVar("_SCT", bound=generic)
+_ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
 
+_DTypeLike = Union[
+    dtype[_SCT],
+    Type[_SCT],
+    _SupportsDType[dtype[_SCT]],
+]
+_ArrayLike = _FiniteNestedSequence[_SupportsArray[dtype[_SCT]]]
 _CorrelateMode = Literal["valid", "same", "full"]
 
+__all__: List[str]
+
 @overload
 def zeros_like(
     a: _ArrayType,
@@ -30,20 +75,61 @@ def zeros_like(
 ) -> _ArrayType: ...
 @overload
 def zeros_like(
-    a: ArrayLike,
-    dtype: DTypeLike = ...,
+    a: _ArrayLike[_SCT],
+    dtype: None = ...,
     order: _OrderKACF = ...,
     subok: bool = ...,
-    shape: Optional[_ShapeLike] = ...,
-) -> ndarray: ...
+    shape: None | _ShapeLike = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def zeros_like(
+    a: object,
+    dtype: None = ...,
+    order: _OrderKACF = ...,
+    subok: bool = ...,
+    shape: None | _ShapeLike= ...,
+) -> NDArray[Any]: ...
+@overload
+def zeros_like(
+    a: Any,
+    dtype: _DTypeLike[_SCT],
+    order: _OrderKACF = ...,
+    subok: bool = ...,
+    shape: None | _ShapeLike= ...,
+) -> NDArray[_SCT]: ...
+@overload
+def zeros_like(
+    a: Any,
+    dtype: DTypeLike,
+    order: _OrderKACF = ...,
+    subok: bool = ...,
+    shape: None | _ShapeLike= ...,
+) -> NDArray[Any]: ...
 
+@overload
 def ones(
     shape: _ShapeLike,
-    dtype: DTypeLike = ...,
+    dtype: None = ...,
+    order: _OrderCF = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[float64]: ...
+@overload
+def ones(
+    shape: _ShapeLike,
+    dtype: _DTypeLike[_SCT],
     order: _OrderCF = ...,
     *,
     like: ArrayLike = ...,
-) -> ndarray: ...
+) -> NDArray[_SCT]: ...
+@overload
+def ones(
+    shape: _ShapeLike,
+    dtype: DTypeLike,
+    order: _OrderCF = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[Any]: ...
 
 @overload
 def ones_like(
@@ -55,21 +141,64 @@ def ones_like(
 ) -> _ArrayType: ...
 @overload
 def ones_like(
-    a: ArrayLike,
-    dtype: DTypeLike = ...,
+    a: _ArrayLike[_SCT],
+    dtype: None = ...,
     order: _OrderKACF = ...,
     subok: bool = ...,
-    shape: Optional[_ShapeLike] = ...,
-) -> ndarray: ...
+    shape: None | _ShapeLike = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def ones_like(
+    a: object,
+    dtype: None = ...,
+    order: _OrderKACF = ...,
+    subok: bool = ...,
+    shape: None | _ShapeLike= ...,
+) -> NDArray[Any]: ...
+@overload
+def ones_like(
+    a: Any,
+    dtype: _DTypeLike[_SCT],
+    order: _OrderKACF = ...,
+    subok: bool = ...,
+    shape: None | _ShapeLike= ...,
+) -> NDArray[_SCT]: ...
+@overload
+def ones_like(
+    a: Any,
+    dtype: DTypeLike,
+    order: _OrderKACF = ...,
+    subok: bool = ...,
+    shape: None | _ShapeLike= ...,
+) -> NDArray[Any]: ...
 
+@overload
 def full(
     shape: _ShapeLike,
     fill_value: Any,
-    dtype: DTypeLike = ...,
+    dtype: None = ...,
+    order: _OrderCF = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[Any]: ...
+@overload
+def full(
+    shape: _ShapeLike,
+    fill_value: Any,
+    dtype: _DTypeLike[_SCT],
     order: _OrderCF = ...,
     *,
     like: ArrayLike = ...,
-) -> ndarray: ...
+) -> NDArray[_SCT]: ...
+@overload
+def full(
+    shape: _ShapeLike,
+    fill_value: Any,
+    dtype: DTypeLike,
+    order: _OrderCF = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[Any]: ...
 
 @overload
 def full_like(
@@ -82,13 +211,40 @@ def full_like(
 ) -> _ArrayType: ...
 @overload
 def full_like(
-    a: ArrayLike,
+    a: _ArrayLike[_SCT],
     fill_value: Any,
-    dtype: DTypeLike = ...,
+    dtype: None = ...,
+    order: _OrderKACF = ...,
+    subok: bool = ...,
+    shape: None | _ShapeLike = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def full_like(
+    a: object,
+    fill_value: Any,
+    dtype: None = ...,
+    order: _OrderKACF = ...,
+    subok: bool = ...,
+    shape: None | _ShapeLike= ...,
+) -> NDArray[Any]: ...
+@overload
+def full_like(
+    a: Any,
+    fill_value: Any,
+    dtype: _DTypeLike[_SCT],
+    order: _OrderKACF = ...,
+    subok: bool = ...,
+    shape: None | _ShapeLike= ...,
+) -> NDArray[_SCT]: ...
+@overload
+def full_like(
+    a: Any,
+    fill_value: Any,
+    dtype: DTypeLike,
     order: _OrderKACF = ...,
     subok: bool = ...,
-    shape: Optional[_ShapeLike] = ...,
-) -> ndarray: ...
+    shape: None | _ShapeLike= ...,
+) -> NDArray[Any]: ...
 
 @overload
 def count_nonzero(
@@ -105,78 +261,306 @@ def count_nonzero(
     keepdims: bool = ...,
 ) -> Any: ...  # TODO: np.intp or ndarray[np.intp]
 
-def isfortran(a: Union[ndarray, generic]) -> bool: ...
+def isfortran(a: NDArray[Any] | generic) -> bool: ...
 
-def argwhere(a: ArrayLike) -> ndarray: ...
+def argwhere(a: ArrayLike) -> NDArray[intp]: ...
 
-def flatnonzero(a: ArrayLike) -> ndarray: ...
+def flatnonzero(a: ArrayLike) -> NDArray[intp]: ...
 
+@overload
 def correlate(
-    a: ArrayLike,
-    v: ArrayLike,
+    a: _ArrayLikeBool_co,
+    v: _ArrayLikeBool_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[bool_]: ...
+@overload
+def correlate(
+    a: _ArrayLikeUInt_co,
+    v: _ArrayLikeUInt_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[unsignedinteger[Any]]: ...
+@overload
+def correlate(
+    a: _ArrayLikeInt_co,
+    v: _ArrayLikeInt_co,
     mode: _CorrelateMode = ...,
-) -> ndarray: ...
+) -> NDArray[signedinteger[Any]]: ...
+@overload
+def correlate(
+    a: _ArrayLikeFloat_co,
+    v: _ArrayLikeFloat_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[floating[Any]]: ...
+@overload
+def correlate(
+    a: _ArrayLikeComplex_co,
+    v: _ArrayLikeComplex_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[complexfloating[Any, Any]]: ...
+@overload
+def correlate(
+    a: _ArrayLikeTD64_co,
+    v: _ArrayLikeTD64_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[timedelta64]: ...
+@overload
+def correlate(
+    a: _ArrayLikeObject_co,
+    v: _ArrayLikeObject_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[object_]: ...
 
+@overload
 def convolve(
-    a: ArrayLike,
-    v: ArrayLike,
+    a: _ArrayLikeBool_co,
+    v: _ArrayLikeBool_co,
     mode: _CorrelateMode = ...,
-) -> ndarray: ...
+) -> NDArray[bool_]: ...
+@overload
+def convolve(
+    a: _ArrayLikeUInt_co,
+    v: _ArrayLikeUInt_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[unsignedinteger[Any]]: ...
+@overload
+def convolve(
+    a: _ArrayLikeInt_co,
+    v: _ArrayLikeInt_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[signedinteger[Any]]: ...
+@overload
+def convolve(
+    a: _ArrayLikeFloat_co,
+    v: _ArrayLikeFloat_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[floating[Any]]: ...
+@overload
+def convolve(
+    a: _ArrayLikeComplex_co,
+    v: _ArrayLikeComplex_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[complexfloating[Any, Any]]: ...
+@overload
+def convolve(
+    a: _ArrayLikeTD64_co,
+    v: _ArrayLikeTD64_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[timedelta64]: ...
+@overload
+def convolve(
+    a: _ArrayLikeObject_co,
+    v: _ArrayLikeObject_co,
+    mode: _CorrelateMode = ...,
+) -> NDArray[object_]: ...
 
 @overload
 def outer(
-    a: ArrayLike,
-    b: ArrayLike,
+    a: _ArrayLikeBool_co,
+    b: _ArrayLikeBool_co,
     out: None = ...,
-) -> ndarray: ...
+) -> NDArray[bool_]: ...
 @overload
 def outer(
-    a: ArrayLike,
-    b: ArrayLike,
-    out: _ArrayType = ...,
+    a: _ArrayLikeUInt_co,
+    b: _ArrayLikeUInt_co,
+    out: None = ...,
+) -> NDArray[unsignedinteger[Any]]: ...
+@overload
+def outer(
+    a: _ArrayLikeInt_co,
+    b: _ArrayLikeInt_co,
+    out: None = ...,
+) -> NDArray[signedinteger[Any]]: ...
+@overload
+def outer(
+    a: _ArrayLikeFloat_co,
+    b: _ArrayLikeFloat_co,
+    out: None = ...,
+) -> NDArray[floating[Any]]: ...
+@overload
+def outer(
+    a: _ArrayLikeComplex_co,
+    b: _ArrayLikeComplex_co,
+    out: None = ...,
+) -> NDArray[complexfloating[Any, Any]]: ...
+@overload
+def outer(
+    a: _ArrayLikeTD64_co,
+    b: _ArrayLikeTD64_co,
+    out: None = ...,
+) -> NDArray[timedelta64]: ...
+@overload
+def outer(
+    a: _ArrayLikeObject_co,
+    b: _ArrayLikeObject_co,
+    out: None = ...,
+) -> NDArray[object_]: ...
+@overload
+def outer(
+    a: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co,
+    b: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co,
+    out: _ArrayType,
 ) -> _ArrayType: ...
 
+@overload
 def tensordot(
-    a: ArrayLike,
-    b: ArrayLike,
-    axes: Union[int, Tuple[_ShapeLike, _ShapeLike]] = ...,
-) -> ndarray: ...
+    a: _ArrayLikeBool_co,
+    b: _ArrayLikeBool_co,
+    axes: int | Tuple[_ShapeLike, _ShapeLike] = ...,
+) -> NDArray[bool_]: ...
+@overload
+def tensordot(
+    a: _ArrayLikeUInt_co,
+    b: _ArrayLikeUInt_co,
+    axes: int | Tuple[_ShapeLike, _ShapeLike] = ...,
+) -> NDArray[unsignedinteger[Any]]: ...
+@overload
+def tensordot(
+    a: _ArrayLikeInt_co,
+    b: _ArrayLikeInt_co,
+    axes: int | Tuple[_ShapeLike, _ShapeLike] = ...,
+) -> NDArray[signedinteger[Any]]: ...
+@overload
+def tensordot(
+    a: _ArrayLikeFloat_co,
+    b: _ArrayLikeFloat_co,
+    axes: int | Tuple[_ShapeLike, _ShapeLike] = ...,
+) -> NDArray[floating[Any]]: ...
+@overload
+def tensordot(
+    a: _ArrayLikeComplex_co,
+    b: _ArrayLikeComplex_co,
+    axes: int | Tuple[_ShapeLike, _ShapeLike] = ...,
+) -> NDArray[complexfloating[Any, Any]]: ...
+@overload
+def tensordot(
+    a: _ArrayLikeTD64_co,
+    b: _ArrayLikeTD64_co,
+    axes: int | Tuple[_ShapeLike, _ShapeLike] = ...,
+) -> NDArray[timedelta64]: ...
+@overload
+def tensordot(
+    a: _ArrayLikeObject_co,
+    b: _ArrayLikeObject_co,
+    axes: int | Tuple[_ShapeLike, _ShapeLike] = ...,
+) -> NDArray[object_]: ...
 
+@overload
+def roll(
+    a: _ArrayLike[_SCT],
+    shift: _ShapeLike,
+    axis: None | _ShapeLike = ...,
+) -> NDArray[_SCT]: ...
+@overload
 def roll(
     a: ArrayLike,
     shift: _ShapeLike,
-    axis: Optional[_ShapeLike] = ...,
-) -> ndarray: ...
+    axis: None | _ShapeLike = ...,
+) -> NDArray[Any]: ...
 
-def rollaxis(a: ndarray, axis: int, start: int = ...) -> ndarray: ...
+def rollaxis(
+    a: NDArray[_SCT],
+    axis: int,
+    start: int = ...,
+) -> NDArray[_SCT]: ...
 
 def moveaxis(
-    a: ndarray,
+    a: NDArray[_SCT],
     source: _ShapeLike,
     destination: _ShapeLike,
-) -> ndarray: ...
+) -> NDArray[_SCT]: ...
 
+@overload
 def cross(
-    a: ArrayLike,
-    b: ArrayLike,
+    a: _ArrayLikeBool_co,
+    b: _ArrayLikeBool_co,
     axisa: int = ...,
     axisb: int = ...,
     axisc: int = ...,
-    axis: Optional[int] = ...,
-) -> ndarray: ...
+    axis: None | int = ...,
+) -> NoReturn: ...
+@overload
+def cross(
+    a: _ArrayLikeUInt_co,
+    b: _ArrayLikeUInt_co,
+    axisa: int = ...,
+    axisb: int = ...,
+    axisc: int = ...,
+    axis: None | int = ...,
+) -> NDArray[unsignedinteger[Any]]: ...
+@overload
+def cross(
+    a: _ArrayLikeInt_co,
+    b: _ArrayLikeInt_co,
+    axisa: int = ...,
+    axisb: int = ...,
+    axisc: int = ...,
+    axis: None | int = ...,
+) -> NDArray[signedinteger[Any]]: ...
+@overload
+def cross(
+    a: _ArrayLikeFloat_co,
+    b: _ArrayLikeFloat_co,
+    axisa: int = ...,
+    axisb: int = ...,
+    axisc: int = ...,
+    axis: None | int = ...,
+) -> NDArray[floating[Any]]: ...
+@overload
+def cross(
+    a: _ArrayLikeComplex_co,
+    b: _ArrayLikeComplex_co,
+    axisa: int = ...,
+    axisb: int = ...,
+    axisc: int = ...,
+    axis: None | int = ...,
+) -> NDArray[complexfloating[Any, Any]]: ...
+@overload
+def cross(
+    a: _ArrayLikeObject_co,
+    b: _ArrayLikeObject_co,
+    axisa: int = ...,
+    axisb: int = ...,
+    axisc: int = ...,
+    axis: None | int = ...,
+) -> NDArray[object_]: ...
 
 @overload
 def indices(
     dimensions: Sequence[int],
-    dtype: DTypeLike = ...,
+    dtype: Type[int] = ...,
     sparse: Literal[False] = ...,
-) -> ndarray: ...
+) -> NDArray[int_]: ...
 @overload
 def indices(
     dimensions: Sequence[int],
-    dtype: DTypeLike = ...,
+    dtype: Type[int] = ...,
     sparse: Literal[True] = ...,
-) -> Tuple[ndarray, ...]: ...
+) -> Tuple[NDArray[int_], ...]: ...
+@overload
+def indices(
+    dimensions: Sequence[int],
+    dtype: _DTypeLike[_SCT],
+    sparse: Literal[False] = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def indices(
+    dimensions: Sequence[int],
+    dtype: _DTypeLike[_SCT],
+    sparse: Literal[True],
+) -> Tuple[NDArray[_SCT], ...]: ...
+@overload
+def indices(
+    dimensions: Sequence[int],
+    dtype: DTypeLike,
+    sparse: Literal[False] = ...,
+) -> NDArray[Any]: ...
+@overload
+def indices(
+    dimensions: Sequence[int],
+    dtype: DTypeLike,
+    sparse: Literal[True],
+) -> Tuple[NDArray[Any], ...]: ...
 
 def fromfunction(
     function: Callable[..., _T],
@@ -187,18 +571,39 @@ def fromfunction(
     **kwargs: Any,
 ) -> _T: ...
 
-def isscalar(element: Any) -> bool: ...
+def isscalar(element: object) -> TypeGuard[
+    generic | bool | int | float | complex | str | bytes | memoryview
+]: ...
 
-def binary_repr(num: int, width: Optional[int] = ...) -> str: ...
+def binary_repr(num: int, width: None | int = ...) -> str: ...
 
-def base_repr(number: int, base: int = ..., padding: int = ...) -> str: ...
+def base_repr(
+    number: SupportsAbs[float],
+    base: float = ...,
+    padding: SupportsIndex = ...,
+) -> str: ...
 
+@overload
 def identity(
     n: int,
-    dtype: DTypeLike = ...,
+    dtype: None = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[float64]: ...
+@overload
+def identity(
+    n: int,
+    dtype: _DTypeLike[_SCT],
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def identity(
+    n: int,
+    dtype: DTypeLike,
     *,
     like: ArrayLike = ...,
-) -> ndarray: ...
+) -> NDArray[Any]: ...
 
 def allclose(
     a: ArrayLike,
@@ -208,13 +613,22 @@ def allclose(
     equal_nan: bool = ...,
 ) -> bool: ...
 
+@overload
+def isclose(
+    a: _ScalarLike_co,
+    b: _ScalarLike_co,
+    rtol: float = ...,
+    atol: float = ...,
+    equal_nan: bool = ...,
+) -> bool_: ...
+@overload
 def isclose(
     a: ArrayLike,
     b: ArrayLike,
     rtol: float = ...,
     atol: float = ...,
     equal_nan: bool = ...,
-) -> Any: ...
+) -> NDArray[bool_]: ...
 
 def array_equal(a1: ArrayLike, a2: ArrayLike, equal_nan: bool = ...) -> bool: ...