summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--numpy/__init__.pyi169
-rw-r--r--numpy/tests/typing/fail/fromnumeric.py35
-rw-r--r--numpy/tests/typing/pass/fromnumeric.py69
-rw-r--r--numpy/tests/typing/reveal/fromnumeric.py70
4 files changed, 343 insertions, 0 deletions
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index 8de9cce02..847b451f7 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -886,6 +886,9 @@ _ScalarGenericDT = TypeVar(
"_ScalarGenericDT", bound=Union[dt.datetime, dt.timedelta, generic]
)
+_Number = TypeVar('_Number', bound=number)
+_NumberLike = Union[int, float, complex, number, bool_]
+
# An array-like object consisting of integers
_Int = Union[int, integer]
_Bool = Union[bool, bool_]
@@ -901,6 +904,11 @@ _ArrayLikeIntOrBool = Union[
Sequence[_ArrayLikeIntNested],
Sequence[_ArrayLikeBoolNested],
]
+_ArrayLikeBool = Union[
+ _Bool,
+ Sequence[_Bool],
+ ndarray
+]
# The signature of take() follows a common theme with its overloads:
# 1. A generic comes in; the same generic comes out
@@ -1062,3 +1070,164 @@ def compress(
axis: Optional[int] = ...,
out: Optional[ndarray] = ...,
) -> ndarray: ...
+@overload
+def clip(
+ a: _Number,
+ a_min: ArrayLike,
+ a_max: Optional[ArrayLike],
+ out: Optional[ndarray] = ...,
+ **kwargs: Any,
+) -> _Number: ...
+@overload
+def clip(
+ a: _Number,
+ a_min: None,
+ a_max: ArrayLike,
+ out: Optional[ndarray] = ...,
+ **kwargs: Any,
+) -> _Number: ...
+@overload
+def clip(
+ a: ArrayLike,
+ a_min: ArrayLike,
+ a_max: Optional[ArrayLike],
+ out: Optional[ndarray] = ...,
+ **kwargs: Any,
+) -> Union[number, ndarray]: ...
+@overload
+def clip(
+ a: ArrayLike,
+ a_min: None,
+ a_max: ArrayLike,
+ out: Optional[ndarray] = ...,
+ **kwargs: Any,
+) -> Union[number, ndarray]: ...
+@overload
+def sum(
+ a: _Number,
+ axis: Optional[_ShapeLike] = ...,
+ dtype: DtypeLike = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: bool = ...,
+ initial: _NumberLike = ...,
+ where: _ArrayLikeBool = ...,
+) -> _Number: ...
+@overload
+def sum(
+ a: ArrayLike,
+ axis: _ShapeLike = ...,
+ dtype: DtypeLike = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: bool = ...,
+ initial: _NumberLike = ...,
+ where: _ArrayLikeBool = ...,
+) -> Union[number, ndarray]: ...
+@overload
+def all(
+ a: ArrayLike,
+ axis: None = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: Literal[False] = ...,
+) -> bool_: ...
+@overload
+def all(
+ a: ArrayLike,
+ axis: Optional[_ShapeLike] = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: bool = ...,
+) -> Union[bool_, ndarray]: ...
+@overload
+def any(
+ a: ArrayLike,
+ axis: None = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: Literal[False] = ...,
+) -> bool_: ...
+@overload
+def any(
+ a: ArrayLike,
+ axis: Optional[_ShapeLike] = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: bool = ...,
+) -> Union[bool_, ndarray]: ...
+def cumsum(
+ a: ArrayLike,
+ axis: Optional[int] = ...,
+ dtype: DtypeLike = ...,
+ out: Optional[ndarray] = ...,
+) -> ndarray: ...
+@overload
+def ptp(
+ a: _Number,
+ axis: Optional[_ShapeLike] = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: bool = ...,
+) -> _Number: ...
+@overload
+def ptp(
+ a: ArrayLike,
+ axis: None = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: Literal[False] = ...,
+) -> number: ...
+@overload
+def ptp(
+ a: ArrayLike,
+ axis: Optional[_ShapeLike] = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: bool = ...,
+) -> Union[number, ndarray]: ...
+@overload
+def amax(
+ a: _Number,
+ axis: Optional[_ShapeLike] = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: bool = ...,
+ initial: _NumberLike = ...,
+ where: _ArrayLikeBool = ...,
+) -> _Number: ...
+@overload
+def amax(
+ a: ArrayLike,
+ axis: None = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: Literal[False] = ...,
+ initial: _NumberLike = ...,
+ where: _ArrayLikeBool = ...,
+) -> number: ...
+@overload
+def amax(
+ a: ArrayLike,
+ axis: Optional[_ShapeLike] = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: bool = ...,
+ initial: _NumberLike = ...,
+ where: _ArrayLikeBool = ...,
+) -> Union[number, ndarray]: ...
+@overload
+def amin(
+ a: _Number,
+ axis: Optional[_ShapeLike] = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: bool = ...,
+ initial: _NumberLike = ...,
+ where: _ArrayLikeBool = ...,
+) -> _Number: ...
+@overload
+def amin(
+ a: ArrayLike,
+ axis: None = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: Literal[False] = ...,
+ initial: _NumberLike = ...,
+ where: _ArrayLikeBool = ...,
+) -> number: ...
+@overload
+def amin(
+ a: ArrayLike,
+ axis: Optional[_ShapeLike] = ...,
+ out: Optional[ndarray] = ...,
+ keepdims: bool = ...,
+ initial: _NumberLike = ...,
+ where: _ArrayLikeBool = ...,
+) -> Union[number, ndarray]: ...
diff --git a/numpy/tests/typing/fail/fromnumeric.py b/numpy/tests/typing/fail/fromnumeric.py
index 7455ce722..66f8a89d0 100644
--- a/numpy/tests/typing/fail/fromnumeric.py
+++ b/numpy/tests/typing/fail/fromnumeric.py
@@ -89,3 +89,38 @@ np.ravel(a, order="bob") # E: Argument "order" to "ravel" has incompatible type
np.compress(
[True], A, axis=1.0 # E: Argument "axis" to "compress" has incompatible type
)
+
+np.clip(a, 1, 2, out=1) # E: No overload variant of "clip" matches argument type
+np.clip(1, None, None) # E: No overload variant of "clip" matches argument type
+
+np.sum(a, axis=1.0) # E: No overload variant of "sum" matches argument type
+np.sum(a, keepdims=1.0) # E: No overload variant of "sum" matches argument type
+np.sum(a, initial=[1]) # E: No overload variant of "sum" matches argument type
+
+np.all(a, axis=1.0) # E: No overload variant of "all" matches argument type
+np.all(a, keepdims=1.0) # E: No overload variant of "all" matches argument type
+np.all(a, out=1.0) # E: No overload variant of "all" matches argument type
+
+np.any(a, axis=1.0) # E: No overload variant of "any" matches argument type
+np.any(a, keepdims=1.0) # E: No overload variant of "any" matches argument type
+np.any(a, out=1.0) # E: No overload variant of "any" matches argument type
+
+np.cumsum(a, axis=1.0) # E: Argument "axis" to "cumsum" has incompatible type
+np.cumsum(a, dtype=1.0) # E: Argument "dtype" to "cumsum" has incompatible type
+np.cumsum(a, out=1.0) # E: Argument "out" to "cumsum" has incompatible type
+
+np.ptp(a, axis=1.0) # E: No overload variant of "ptp" matches argument type
+np.ptp(a, keepdims=1.0) # E: No overload variant of "ptp" matches argument type
+np.ptp(a, out=1.0) # E: No overload variant of "ptp" matches argument type
+
+np.amax(a, axis=1.0) # E: No overload variant of "amax" matches argument type
+np.amax(a, keepdims=1.0) # E: No overload variant of "amax" matches argument type
+np.amax(a, out=1.0) # E: No overload variant of "amax" matches argument type
+np.amax(a, initial=[1.0]) # E: No overload variant of "amax" matches argument type
+np.amax(a, where=[1.0]) # E: List item 0 has incompatible type
+
+np.amin(a, axis=1.0) # E: No overload variant of "amin" matches argument type
+np.amin(a, keepdims=1.0) # E: No overload variant of "amin" matches argument type
+np.amin(a, out=1.0) # E: No overload variant of "amin" matches argument type
+np.amin(a, initial=[1.0]) # E: No overload variant of "amin" matches argument type
+np.amin(a, where=[1.0]) # E: List item 0 has incompatible type
diff --git a/numpy/tests/typing/pass/fromnumeric.py b/numpy/tests/typing/pass/fromnumeric.py
index 0ce8ef970..d9dd45c54 100644
--- a/numpy/tests/typing/pass/fromnumeric.py
+++ b/numpy/tests/typing/pass/fromnumeric.py
@@ -114,3 +114,72 @@ np.compress([True], b)
np.compress([True], c)
np.compress([True], A)
np.compress([True], B)
+
+np.clip(a, 0, 1.0)
+np.clip(b, -1, 1)
+np.clip(a, 0, None)
+np.clip(b, None, 1)
+np.clip(c, 0, 1)
+np.clip(A, 0, 1)
+np.clip(B, 0, 1)
+np.clip(B, [0, 1], [1, 2])
+
+np.sum(a)
+np.sum(b)
+np.sum(c)
+np.sum(A)
+np.sum(B)
+np.sum(A, axis=0)
+np.sum(B, axis=0)
+
+np.all(a)
+np.all(b)
+np.all(c)
+np.all(A)
+np.all(B)
+np.all(A, axis=0)
+np.all(B, axis=0)
+np.all(A, keepdims=True)
+np.all(B, keepdims=True)
+
+np.any(a)
+np.any(b)
+np.any(c)
+np.any(A)
+np.any(B)
+np.any(A, axis=0)
+np.any(B, axis=0)
+np.any(A, keepdims=True)
+np.any(B, keepdims=True)
+
+np.cumsum(a)
+np.cumsum(b)
+np.cumsum(c)
+np.cumsum(A)
+np.cumsum(B)
+
+np.ptp(b)
+np.ptp(c)
+np.ptp(B)
+np.ptp(B, axis=0)
+np.ptp(B, keepdims=True)
+
+np.amax(a)
+np.amax(b)
+np.amax(c)
+np.amax(A)
+np.amax(B)
+np.amax(A, axis=0)
+np.amax(B, axis=0)
+np.amax(A, keepdims=True)
+np.amax(B, keepdims=True)
+
+np.amin(a)
+np.amin(b)
+np.amin(c)
+np.amin(A)
+np.amin(B)
+np.amin(A, axis=0)
+np.amin(B, axis=0)
+np.amin(A, keepdims=True)
+np.amin(B, keepdims=True)
diff --git a/numpy/tests/typing/reveal/fromnumeric.py b/numpy/tests/typing/reveal/fromnumeric.py
index 7d79d5fa9..f5feb3f5f 100644
--- a/numpy/tests/typing/reveal/fromnumeric.py
+++ b/numpy/tests/typing/reveal/fromnumeric.py
@@ -133,3 +133,73 @@ reveal_type(np.compress([True], b)) # E: numpy.ndarray
reveal_type(np.compress([True], c)) # E: numpy.ndarray
reveal_type(np.compress([True], A)) # E: numpy.ndarray
reveal_type(np.compress([True], B)) # E: numpy.ndarray
+
+reveal_type(np.clip(a, 0, 1.0)) # E: numpy.number
+reveal_type(np.clip(b, -1, 1)) # E: numpy.float32
+reveal_type(np.clip(c, 0, 1)) # E: numpy.number
+reveal_type(np.clip(A, 0, 1)) # E: Union[numpy.number, numpy.ndarray]
+reveal_type(np.clip(B, 0, 1)) # E: Union[numpy.number, numpy.ndarray]
+
+reveal_type(np.sum(a)) # E: numpy.number
+reveal_type(np.sum(b)) # E: numpy.float32
+reveal_type(np.sum(c)) # E: numpy.number
+reveal_type(np.sum(A)) # E: numpy.number
+reveal_type(np.sum(B)) # E: numpy.number
+reveal_type(np.sum(A, axis=0)) # E: Union[numpy.number, numpy.ndarray]
+reveal_type(np.sum(B, axis=0)) # E: Union[numpy.number, numpy.ndarray]
+
+reveal_type(np.all(a)) # E: numpy.bool_
+reveal_type(np.all(b)) # E: numpy.bool_
+reveal_type(np.all(c)) # E: numpy.bool_
+reveal_type(np.all(A)) # E: numpy.bool_
+reveal_type(np.all(B)) # E: numpy.bool_
+reveal_type(np.all(A, axis=0)) # E: Union[numpy.bool_, numpy.ndarray]
+reveal_type(np.all(B, axis=0)) # E: Union[numpy.bool_, numpy.ndarray]
+reveal_type(np.all(A, keepdims=True)) # E: Union[numpy.bool_, numpy.ndarray]
+reveal_type(np.all(B, keepdims=True)) # E: Union[numpy.bool_, numpy.ndarray]
+
+reveal_type(np.any(a)) # E: numpy.bool_
+reveal_type(np.any(b)) # E: numpy.bool_
+reveal_type(np.any(c)) # E: numpy.bool_
+reveal_type(np.any(A)) # E: numpy.bool_
+reveal_type(np.any(B)) # E: numpy.bool_
+reveal_type(np.any(A, axis=0)) # E: Union[numpy.bool_, numpy.ndarray]
+reveal_type(np.any(B, axis=0)) # E: Union[numpy.bool_, numpy.ndarray]
+reveal_type(np.any(A, keepdims=True)) # E: Union[numpy.bool_, numpy.ndarray]
+reveal_type(np.any(B, keepdims=True)) # E: Union[numpy.bool_, numpy.ndarray]
+
+reveal_type(np.cumsum(a)) # E: numpy.ndarray
+reveal_type(np.cumsum(b)) # E: numpy.ndarray
+reveal_type(np.cumsum(c)) # E: numpy.ndarray
+reveal_type(np.cumsum(A)) # E: numpy.ndarray
+reveal_type(np.cumsum(B)) # E: numpy.ndarray
+
+reveal_type(np.ptp(a)) # E: numpy.number
+reveal_type(np.ptp(b)) # E: numpy.float32
+reveal_type(np.ptp(c)) # E: numpy.number
+reveal_type(np.ptp(A)) # E: numpy.number
+reveal_type(np.ptp(B)) # E: numpy.number
+reveal_type(np.ptp(A, axis=0)) # E: Union[numpy.number, numpy.ndarray]
+reveal_type(np.ptp(B, axis=0)) # E: Union[numpy.number, numpy.ndarray]
+reveal_type(np.ptp(A, keepdims=True)) # E: Union[numpy.number, numpy.ndarray]
+reveal_type(np.ptp(B, keepdims=True)) # E: Union[numpy.number, numpy.ndarray]
+
+reveal_type(np.amax(a)) # E: numpy.number
+reveal_type(np.amax(b)) # E: numpy.float32
+reveal_type(np.amax(c)) # E: numpy.number
+reveal_type(np.amax(A)) # E: numpy.number
+reveal_type(np.amax(B)) # E: numpy.number
+reveal_type(np.amax(A, axis=0)) # E: Union[numpy.number, numpy.ndarray]
+reveal_type(np.amax(B, axis=0)) # E: Union[numpy.number, numpy.ndarray]
+reveal_type(np.amax(A, keepdims=True)) # E: Union[numpy.number, numpy.ndarray]
+reveal_type(np.amax(B, keepdims=True)) # E: Union[numpy.number, numpy.ndarray]
+
+reveal_type(np.amin(a)) # E: numpy.number
+reveal_type(np.amin(b)) # E: numpy.float32
+reveal_type(np.amin(c)) # E: numpy.number
+reveal_type(np.amin(A)) # E: numpy.number
+reveal_type(np.amin(B)) # E: numpy.number
+reveal_type(np.amin(A, axis=0)) # E: Union[numpy.number, numpy.ndarray]
+reveal_type(np.amin(B, axis=0)) # E: Union[numpy.number, numpy.ndarray]
+reveal_type(np.amin(A, keepdims=True)) # E: Union[numpy.number, numpy.ndarray]
+reveal_type(np.amin(B, keepdims=True)) # E: Union[numpy.number, numpy.ndarray]
View Lorry Controller Status