8000 ENH: Add type annotations for the np.core.fromnumeric module: part 2/… · numpy/numpy-stubs@5aaa332 · GitHub
[go: up one dir, main page]
Skip to content
This repository was archived by the owner on Jun 10, 2020. It is now read-only.

Commit 5aaa332

Browse files
BvB93BvB93
authored
ENH: Add type annotations for the np.core.fromnumeric module: part 2/4 (#71)
* Added 11 new functions + tests and updated some existing annotations New functions: * ``argmax()`` * ``argmin()`` * ``searchsorted()`` * ``resize()`` * ``squeeze()`` * ``diagonal()`` * ``trace()`` * ``ravel()`` * ``nonzero()`` * ``shape()`` * ``compress()``
1 parent 23e88e0 commit 5aaa332

File tree

4 files changed

+220
-2
lines changed

4 files changed

+220
-2
lines changed

‎numpy-stubs/__init__.pyi

Lines changed: 69 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -801,6 +801,7 @@ _Mode = Literal["raise", "wrap", "clip"]
801801
_Order = Literal["C", "F", "A"]
802802
_PartitionKind = Literal["introselect"]
803803
_SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"]
804+
_Side = Literal["left", "right"]
804805

805806
# Various annotations for scalars
806807

@@ -822,8 +823,8 @@ _ScalarGenericDT = TypeVar(
822823
_Int = Union[int, integer]
823824
_Bool = Union[bool, bool_]
824825
_IntOrBool = Union[_Int, _Bool]
825-
_ArrayLikeIntNested = Any # TODO: wait for support for recursive types
826-
_ArrayLikeBoolNested = Any # TODO: wait for support for recursive types
826+
_ArrayLikeIntNested = ArrayLike # TODO: wait for support for recursive types
827+
_ArrayLikeBoolNested = ArrayLike # TODO: wait for support for recursive types
827828

828829
# Integers and booleans can generally be used interchangeably
829830
_ArrayLikeIntOrBool = Union[
@@ -948,3 +949,69 @@ def argsort(
948949
kind: Optional[_SortKind] = ...,
949950
order: Union[None, str, Sequence[str]] = ...,
950951
) -> ndarray: ...
952+
@overload
953+
def argmax(
954+
a: Union[Sequence[ArrayLike], ndarray],
955+
axis: None = ...,
956+
out: Optional[ndarray] = ...,
957+
) -> integer: ...
958+
@overload
959+
def argmax(
960+
a: Union[Sequence[ArrayLike], ndarray],
961+
axis: int = ...,
962+
out: Optional[ndarray] = ...,
963+
) -> Union[integer, ndarray]: ...
964+
@overload
965+
def argmin(
966+
a: Union[Sequence[ArrayLike], ndarray],
967+
axis: None = ...,
968+
out: Optional[ndarray] = ...,
969+
) -> integer: ...
970+
@overload
971+
def argmin(
972+
a: Union[Sequence[ArrayLike], ndarray],
973+
axis: int = ...,
974+
out: Optional[ndarray] = ...,
975+
) -> Union[integer, ndarray]: ...
976+
@overload
977+
def searchsorted(
978+
a: Union[Sequence[ArrayLike], ndarray],
979+
v: _Scalar,
980+
side: _Side = ...,
981+
sorter: Union[None, Sequence[_IntOrBool], ndarray] = ..., # 1D int array
982+
) -> integer: ...
983+
@overload
984+
def searchsorted(
985+
a: Union[Sequence[ArrayLike], ndarray],
986+
v: ArrayLike,
987+
side: _Side = ...,
57AE 988+
sorter: Union[None, Sequence[_IntOrBool], ndarray] = ..., # 1D int array
989+
) -> ndarray: ...
990+
def resize(a: ArrayLike, new_shape: _ShapeLike) -> ndarray: ...
991+
@overload
992+
def squeeze(a: _ScalarGeneric, axis: Optional[_ShapeLike] = ...) -> _ScalarGeneric: ...
993+
@overload
994+
def squeeze(a: ArrayLike, axis: Optional[_ShapeLike] = ...) -> ndarray: ...
995+
def diagonal(
996+
a: Union[Sequence[Sequence[ArrayLike]], ndarray], # >= 2D array
997+
offset: int = ...,
998+
axis1: int = ...,
999+
axis2: int = ...,
1000+
) -> ndarray: ...
1001+
def trace(
1002+
a: Union[Sequence[Sequence[ArrayLike]], ndarray], # >= 2D array
1003+
offset: int = ...,
1004+
axis1: int = ...,
1005+
axis2: int = ...,
1006+
dtype: DtypeLike = ...,
1007+
out: Optional[ndarray] = ...,
1008+
) -> Union[number, ndarray]: ...
1009+
def ravel(a: ArrayLike, order: _Order = ...) -> ndarray: ...
1010+
def nonzero(a: ArrayLike) -> Tuple[ndarray, ...]: ...
1011+
def shape(a: ArrayLike) -> _Shape: ...
1012+
def compress(
1013+
condition: Union[Sequence[_Bool], ndarray], # 1D bool array
1014+
a: ArrayLike,
1015+
axis: Optional[int] = ...,
1016+
out: Optional[ndarray] = ...,
1017+
) -> ndarray: ...

‎tests/fail/fromnumeric.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,3 +62,40 @@
6262
np.argsort(A, axis="bob") # E: Argument "axis" to "argsort" has incompatible type
6363
np.argsort(A, kind="bob") # E: Argument "kind" to "argsort" has incompatible type
6464
np.argsort(A, order=range(5)) # E: Argument "order" to "argsort" has incompatible type
65+
66+
np.argmax(a) # E: No overload variant of "argmax" matches argument type
67+
np.argmax(A, axis="bob") # E: No overload variant of "argmax" matches argument type
68+
np.argmax(A, kind="bob") # E: No overload variant of "argmax" matches argument type
69+
70+
np.argmin(a) # E: No overload variant of "argmin" matches argument type
71+
np.argmin(A, axis="bob") # E: No overload variant of "argmin" matches argument typ F438 e
72+
np.argmin(A, kind="bob") # E: No overload variant of "argmin" matches argument type
73+
74+
np.searchsorted(a, 0) # E: No overload variant of "searchsorted" matches argument type
75+
np.searchsorted( # E: No overload variant of "searchsorted" matches argument type
76+
A[0], 0, side="bob"
77+
)
78+
np.searchsorted( # E: No overload variant of "searchsorted" matches argument type
79+
A[0], 0, sorter=1.0
80+
)
81+
82+
np.resize(A, 1.0) # E: Argument 2 to "resize" has incompatible type
83+
84+
np.squeeze(A, 1.0) # E: No overload variant of "squeeze" matches argument type
85+
86+
np.diagonal(a) # E: Argument 1 to "diagonal" has incompatible type
87+
np.diagonal(A, offset=None) # E: Argument "offset" to "diagonal" has incompatible type
88+
np.diagonal(A, axis1="bob") # E: Argument "axis1" to "diagonal" has incompatible type
89+
np.diagonal(A, axis2=[]) # E: Argument "axis2" to "diagonal" has incompatible type
90+
91+
np.trace(a) # E: Argument 1 to "trace" has incompatible type
92+
np.trace(A, offset=None) # E: Argument "offset" to "trace" has incompatible type
93+
np.trace(A, axis1="bob") # E: Argument "axis1" to "trace" has incompatible type
94+
np.trace(A, axis2=[]) # E: Argument "axis2" to "trace" has incompatible type
95+
96+
np.ravel(a, order="bob") # E: Argument "order" to "ravel" has incompatible type
97+
98+
np.compress(True, A) # E: Argument 1 to "compress" has incompatible type
99+
np.compress(
100+
[True], A, axis=1.0 # E: Argument "axis" to "compress" has incompatible type
101+
)

‎tests/pass/fromnumeric.py

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,3 +60,60 @@
6060

6161
np.argsort(A, 0)
6262
np.argsort(B, 0)
63+
64+
np.argmax(A)
65+
np.argmax(B)
66+
np.argmax(A, axis=0)
67+
np.argmax(B, axis=0)
68+
69+
np.argmin(A)
70+
np.argmin(B)
71+
np.argmin(A, axis=0)
72+
np.argmin(B, axis=0)
73+
74+
np.searchsorted(A[0], 0)
75+
np.searchsorted(B[0], 0)
76+
np.searchsorted(A[0], [0])
77+
np.searchsorted(B[0], [0])
78+
79+
np.resize(a, (5, 5))
80+
np.resize(b, (5, 5))
81+
np.resize(c, (5, 5))
82+
np.resize(A, (5, 5))
83+
np.resize(B, (5, 5))
84+
85+
np.squeeze(a)
86+
np.squeeze(b)
87+
np.squeeze(c)
88+
np.squeeze(A)
89+
np.squeeze(B)
90+
91+
np.diagonal(A)
92+
np.diagonal(B)
93+
94+
np.trace(A)
95+
np.trace(B)
96+
97+
np.ravel(a)
98+
np.ravel(b)
99+
np.ravel(c)
100+
np.ravel(A)
101+
np.ravel(B)
102+
103+
np.nonzero(a)
104+
np.nonzero(b)
105+
np.nonzero(c)
106+
np.nonzero(A)
107+
np.nonzero(B)
108+
109+
np.shape(a)
110+
np.shape(b)
111+
np.shape(c)
112+
np.shape(A)
113+
np.shape(B)
114+
115+
np.compress([True], a)
116+
np.compress([True], b)
117+
np.compress([True], c)
118+
np.compress([True], A)
119+
np.compress([True], B)

‎tests/reveal/fromnumeric.py

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,3 +76,60 @@
7676

7777
reveal_type(np.argsort(A, 0)) # E: numpy.ndarray
7878
reveal_type(np.argsort(B, 0)) # E: numpy.ndarray
79+
80+
reveal_type(np.argmax(A)) # E: numpy.integer
81+
reveal_type(np.argmax(B)) # E: numpy.integer
82+
reveal_type(np.argmax(A, axis=0)) # E: Union[numpy.integer, numpy.ndarray]
83+
reveal_type(np.argmax(B, axis=0)) # E: Union[numpy.integer, numpy.ndarray]
84+
85+
reveal_type(np.argmin(A)) # E: numpy.integer
86+
reveal_type(np.argmin(B)) # E: numpy.integer
87+
reveal_type(np.argmin(A, axis=0)) # E: Union[numpy.integer, numpy.ndarray]
88+
reveal_type(np.argmin(B, axis=0)) # E: Union[numpy.integer, numpy.ndarray]
89+
90+
reveal_type(np.searchsorted(A[0], 0)) # E: numpy.integer
91+
reveal_type(np.searchsorted(B[0], 0)) # E: numpy.integer
92+
reveal_type(np.searchsorted(A[0], [0])) # E: numpy.ndarray
93+
reveal_type(np.searchsorted(B[0], [0])) # E: numpy.ndarray
94+
95+
reveal_type(np.resize(a, (5, 5))) # E: numpy.ndarray
96+
reveal_type(np.resize(b, (5, 5))) # E: numpy.ndarray
97+
reveal_type(np.resize(c, (5, 5))) # E: numpy.ndarray
98+
reveal_type(np.resize(A, (5, 5))) # E: numpy.ndarray
99+
reveal_type(np.resize(B, (5, 5))) # E: numpy.ndarray
100+
101+
reveal_type(np.squeeze(a)) # E: numpy.bool_
102+
reveal_type(np.squeeze(b)) # E: numpy.float32
103+
reveal_type(np.squeeze(c)) # E: numpy.ndarray
104+
reveal_type(np.squeeze(A)) # E: numpy.ndarray
105+
reveal_type(np.squeeze(B)) # E: numpy.ndarray
106+
107+
reveal_type(np.diagonal(A)) # E: numpy.ndarray
108+
reveal_type(np.diagonal(B)) # E: numpy.ndarray
109+
110+
reveal_type(np.trace(A)) # E: Union[numpy.number, numpy.ndarray]
111+
reveal_type(np.trace(B)) # E: Union[numpy.number, numpy.ndarray]
112+
113+
reveal_type(np.ravel(a)) # E: numpy.ndarray
114+
reveal_type(np.ravel(b)) # E: numpy.ndarray
115+
reveal_type(np.ravel(c)) # E: numpy.ndarray
116+
reveal_type(np.ravel(A)) # E: numpy.ndarray
117+
reveal_type(np.ravel(B)) # E: numpy.ndarray
118+
119+
reveal_type(np.nonzero(a)) # E: tuple[numpy.ndarray]
120+
reveal_type(np.nonzero(b)) # E: tuple[numpy.ndarray]
121+
reveal_type(np.nonzero(c)) # E: tuple[numpy.ndarray]
122+
reveal_type(np.nonzero(A)) # E: tuple[numpy.ndarray]
123+
reveal_type(np.nonzero(B)) # E: tuple[numpy.ndarray]
124+
125+
reveal_type(np.shape(a)) # E: tuple[builtins.int]
126+
reveal_type(np.shape(b)) # E: tuple[builtins.int]
127+
reveal_type(np.shape(c)) # E: tuple[builtins.int]
128+
reveal_type(np.shape(A)) # E: tuple[builtins.int]
129+
reveal_type(np.shape(B)) # E: tuple[builtins.int]
130+
131+
reveal_type(np.compress([True], a)) # E: numpy.ndarray
132+
reveal_type(np.compress([True], b)) # E: numpy.ndarray
133+
reveal_type(np.compress([True], c)) # E: numpy.ndarray
134+
reveal_type(np.compress([True], A)) # E: numpy.ndarray
135+
reveal_type(np.compress([True], B)) # E: numpy.ndarray

0 commit comments

Comments
 (0)
0