Skip to content
This repository was archived by the owner on Jun 10, 2020. It is now read-only.

Commit 2011153

Browse files
BvB93BvB93
authored
ENH: Add type annotations for the np.core.fromnumeric module: part 1/4 (#67)
Added 11 functions from np.core.fromnumeric: - take - reshape - choose - repeat - put - swapaxes - transpose - partition - argpartition - sort - argsort
1 parent 26757ab commit 2011153

File tree

5 files changed

+340
-0
lines changed

5 files changed

+340
-0
lines changed

‎.gitignore

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,3 +6,6 @@ venv
66
.idea
77
*~
88
**~
9+
10+
# MacOS
11+
.DS_Store

‎numpy-stubs/__init__.pyi

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,11 @@ if sys.version_info[0] < 3:
3535
else:
3636
from typing import SupportsBytes
3737

38+
if sys.version_info >= (3, 8):
39+
from typing import Literal
40+
else:
41+
from typing_extensions import Literal
42+
3843
# TODO: remove when the full numpy namespace is defined
3944
def __getattr__(name: str) -> Any: ...
4045

@@ -792,3 +797,123 @@ class AxisError(ValueError, IndexError):
792797
def __init__(
793798
self, axis: int, ndim: Optional[int] = ..., msg_prefix: Optional[str] = ...
794799
) -> None: ...
800+
801+
# Functions from np.core.fromnumeric
802+
_Mode = Literal["raise", "wrap", "clip"]
803+
_Order = Literal["C", "F", "A"]
804+
_PartitionKind = Literal["introselect"]
805+
_SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"]
806+
807+
# Various annotations for scalars
808+
809+
# While dt.datetime and dt.timedelta are not technically part of NumPy,
810+
# they are one of the rare few builtin scalars which serve as valid return types.
811+
# See https://github.com/numpy/numpy-stubs/pull/67#discussion_r412604113.
812+
_ScalarNumpy = Union[generic, dt.datetime, dt.timedelta]
813+
_ScalarBuiltin = Union[str, bytes, dt.date, dt.timedelta, bool, int, float, complex]
814+
_Scalar = Union[_ScalarBuiltin, _ScalarNumpy]
815+
816+
_ScalarGeneric = TypeVar(
817+
"_ScalarGeneric", bound=Union[dt.datetime, dt.timedelta, generic]
818+
)
819+
820+
# An array-like object consisting of integers
821+
_Int = Union[int, integer]
822+
_ArrayLikeIntNested = Any # TODO: wait for support for recursive types
823+
_ArrayLikeInt = Union[_Int, ndarray, Sequence[_Int], Sequence[_ArrayLikeIntNested]]
824+
825+
# The signature of take() follows a common theme with its overloads:
826+
# 1. A generic comes in; the same generic comes out
827+
# 2. A scalar comes in; a generic comes out
828+
# 3. An array-like object comes in; some keyword ensures that a generic comes out
829+
# 4. An array-like object comes in; an ndarray or generic comes out
830+
@overload
831+
def take(
832+
a: _ScalarGeneric,
833+
indices: int,
834+
axis: Optional[int] = ...,
835+
out: Optional[ndarray] = ...,
836+
mode: _Mode = ...,
837+
) -> _ScalarGeneric: ...
838+
@overload
839+
def take(
840+
a: _Scalar,
841+
indices: int,
842+
axis: Optional[int] = ...,
843+
out: Optional[ndarray] = ...,
844+
mode: _Mode = ...,
845+
) -> _ScalarNumpy: ...
846+
@overload
847+
def take(
848+
a: _ArrayLike,
849+
indices: int,
850+
axis: Optional[int] = ...,
851+
out: Optional[ndarray] = ...,
852+
mode: _Mode = ...,
853+
) -> _ScalarNumpy: ...
854+
@overload
855+
def take(
856+
a: _ArrayLike,
857+
indices: _ArrayLikeInt,
858+
axis: Optional[int] = ...,
859+
out: Optional[ndarray] = ...,
860+
mode: _Mode = ...,
861+
) -> Union[_ScalarNumpy, ndarray]: ...
862+
def reshape(a: _ArrayLike, newshape: _ShapeLike, order: _Order = ...) -> ndarray: ...
863+
@overload
864+
def choose(
865+
a: _ScalarGeneric,
866+
choices: Union[Sequence[_ArrayLike], ndarray],
867+
out: Optional[ndarray] = ...,
868+
mode: _Mode = ...,
869+
) -> _ScalarGeneric: ...
870+
@overload
871+
def choose(
872+
a: _Scalar,
873+
choices: Union[Sequence[_ArrayLike], ndarray],
874+
out: Optional[ndarray] = ...,
875+
mode: _Mode = ...,
876+
) -> _ScalarNumpy: ...
877+
@overload
878+
def choose(
879+
a: _ArrayLike,
880+
choices: Union[Sequence[_ArrayLike], ndarray],
881+
out: Optional[ndarray] = ...,
882+
mode: _Mode = ...,
883+
) -> ndarray: ...
884+
def repeat(
885+
a: _ArrayLike, repeats: _ArrayLikeInt, axis: Optional[int] = ...
886+
) -> ndarray: ...
887+
def put(a: ndarray, ind: _ArrayLikeInt, v: _ArrayLike, mode: _Mode = ...) -> None: ...
888+
def swapaxes(
889+
a: Union[Sequence[_ArrayLike], ndarray], axis1: int, axis2: int
890+
) -> ndarray: ...
891+
def transpose(
892+
a: _ArrayLike, axes: Union[None, Sequence[int], ndarray] = ...
893+
) -> ndarray: ...
894+
def partition(
895+
a: _ArrayLike,
896+
kth: _ArrayLikeInt,
897+
axis: Optional[int] = ...,
898+
kind: _PartitionKind = ...,
899+
order: Union[None, str, Sequence[str]] = ...,
900+
) -> ndarray: ...
901+
def argpartition(
902+
a: _ArrayLike,
903+
kth: _ArrayLikeInt,
904+
axis: Optional[int] = ...,
905+
kind: _PartitionKind = ...,
906+
order: Union[None, str, Sequence[str]] = ...,
907+
) -> ndarray: ...
908+
def sort(
909+
a: Union[Sequence[_ArrayLike], ndarray],
910+
axis: Optional[int] = ...,
911+
kind: Optional[_SortKind] = ...,
912+
order: Union[None, str, Sequence[str]] = ...,
913+
) -> ndarray: ...
914+
def argsort(
915+
a: Union[Sequence[_ArrayLike], ndarray],
916+
axis: Optional[int] = ...,
917+
kind: Optional[_SortKind] = ...,
918+
order: Union[None, str, Sequence[str]] = ...,
919+
) -> ndarray: ...

‎tests/fail/fromnumeric.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
"""Tests for :mod:`numpy.core.fromnumeric`."""
2+
3+
import numpy as np
4+
5+
A = np.array(True, ndmin=2, dtype=bool)
6+
A.setflags(write=False)
7+
8+
a = np.bool_(True)
9+
10+
np.take(a, None) # E: No overload variant of "take" matches argument types
11+
np.take(a, axis=1.0) # E: No overload variant of "take" matches argument types
12+
np.take(A, out=1) # E: No overload variant of "take" matches argument types
13+
np.take(A, mode="bob") # E: No overload variant of "take" matches argument types
14+
15+
np.reshape(a, None) # E: Argument 2 to "reshape" has incompatible type
16+
np.reshape(A, 1, order="bob") # E: Argument "order" to "reshape" has incompatible type
17+
18+
np.choose(a, None) # E: No overload variant of "choose" matches argument types
19+
np.choose(a, out=1.0) # E: No overload variant of "choose" matches argument types
20+
np.choose(A, mode="bob") # E: No overload variant of "choose" matches argument types
21+
22+
np.repeat(a, None) # E: Argument 2 to "repeat" has incompatible type
23+
np.repeat(A, 1, axis=1.0) # E: Argument "axis" to "repeat" has incompatible type
24+
25+
np.swapaxes(a, 0, 0) # E: Argument 1 to "swapaxes" has incompatible type
26+
np.swapaxes(A, None, 1) # E: Argument 2 to "swapaxes" has incompatible type
27+
np.swapaxes(A, 1, [0]) # E: Argument 3 to "swapaxes" has incompatible type
28+
29+
np.transpose(a, axes=1) # E: Argument "axes" to "transpose" has incompatible type
30+
np.transpose(A, axes=1.0) # E: Argument "axes" to "transpose" has incompatible type
31+
32+
np.partition(a, None) # E: Argument 2 to "partition" has incompatible type
33+
np.partition(
34+
a, 0, axis="bob" # E: Argument "axis" to "partition" has incompatible type
35+
)
36+
np.partition(
37+
A, 0, kind="bob" # E: Argument "kind" to "partition" has incompatible type
38+
)
39+
np.partition(
40+
A, 0, order=range(5) # E: Argument "order" to "partition" has incompatible type
41+
)
42+
43+
np.argpartition(a, None) # E: Argument 2 to "argpartition" has incompatible type
44+
np.argpartition(
45+
a, 0, axis="bob" # E: Argument "axis" to "argpartition" has incompatible type
46+
)
47+
np.argpartition(
48+
A, 0, kind="bob" # E: Argument "kind" to "argpartition" has incompatible type
49+
)
50+
np.argpartition(
51+
A, 0, order=range(5) # E: Argument "order" to "argpartition" has incompatible type
52+
)
53+
54+
np.sort(a) # E: Argument 1 to "sort" has incompatible type
55+
np.sort(A, axis="bob") # E: Argument "axis" to "sort" has incompatible type
56+
np.sort(A, kind="bob") # E: Argument "kind" to "sort" has incompatible type
57+
np.sort(A, order=range(5)) # E: Argument "order" to "sort" has incompatible type
58+
59+
np.argsort(a) # E: Argument 1 to "argsort" has incompatible type
60+
np.argsort(A, axis="bob") # E: Argument "axis" to "argsort" has incompatible type
61+
np.argsort(A, kind="bob") # E: Argument "kind" to "argsort" has incompatible type
62+
np.argsort(A, order=range(5)) # E: Argument "order" to "argsort" has incompatible type

‎tests/pass/fromnumeric.py

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
"""Tests for :mod:`numpy.core.fromnumeric`."""
2+
3+
import numpy as np
4+
5+
A = np.array(True, ndmin=2, dtype=bool)
6+
B = np.array(1.0, ndmin=2, dtype=np.float32)
7+
A.setflags(write=False)
8+
B.setflags(write=False)
9+
10+
a = np.bool_(True)
11+
b = np.float32(1.0)
12+
c = 1.0
13+
14+
np.take(a, 0)
15+
np.take(b, 0)
16+
np.take(c, 0)
17+
np.take(A, 0)
18+
np.take(B, 0)
19+
np.take(A, [0])
20+
np.take(B, [0])
21+
22+
np.reshape(a, 1)
23+
np.reshape(b, 1)
24+
np.reshape(c, 1)
25+
np.reshape(A, 1)
26+
np.reshape(B, 1)
27+
28+
np.choose(a, [True])
29+
np.choose(b, [1.0])
30+
np.choose(c, [1.0])
31+
np.choose(A, [True])
32+
np.choose(B, [1.0])
33+
34+
np.repeat(a, 1)
35+
np.repeat(b, 1)
36+
np.repeat(c, 1)
37+
np.repeat(A, 1)
38+
np.repeat(B, 1)
39+
40+
np.swapaxes(A, 0, 0)
41+
np.swapaxes(B, 0, 0)
42+
43+
np.transpose(a)
44+
np.transpose(b)
45+
np.transpose(c)
46+
np.transpose(A)
47+
np.transpose(B)
48+
49+
np.partition(a, 0)
50+
np.partition(b, 0)
51+
np.partition(c, 0)
52+
np.partition(A, 0)
53+
np.partition(B, 0)
54+
55+
np.argpartition(a, 0)
56+
np.argpartition(b, 0)
57+
np.argpartition(c, 0)
58+
np.argpartition(A, 0)
59+
np.argpartition(B, 0)
60+
61+
np.sort(A, 0)
62+
np.sort(B, 0)
63+
64+
np.argsort(A, 0)
65+
np.argsort(B, 0)

‎tests/reveal/fromnumeric.py

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
"""Tests for :mod:`numpy.core.fromnumeric`."""
2+
3+
import numpy as np
4+
5+
A = np.array(True, ndmin=2, dtype=bool)
6+
B = np.array(1.0, ndmin=2, dtype=np.float32)
7+
A.setflags(write=False)
8+
B.setflags(write=False)
9+
10+
a = np.bool_(True)
11+
b = np.float32(1.0)
12+
c = 1.0
13+
14+
reveal_type(np.take(a, 0)) # E: numpy.bool_
15+
reveal_type(np.take(b, 0)) # E: numpy.float32
16+
reveal_type(
17+
np.take(c, 0) # E: Union[numpy.generic, datetime.datetime, datetime.timedelta]
18+
)
19+
reveal_type(
20+
np.take(A, 0) # E: Union[numpy.generic, datetime.datetime, datetime.timedelta]
21+
)
22+
reveal_type(
23+
np.take(B, 0) # E: Union[numpy.generic, datetime.datetime, datetime.timedelta]
24+
)
25+
reveal_type(
26+
np.take( # E: Union[numpy.generic, datetime.datetime, datetime.timedelta, numpy.ndarray]
27+
A, [0]
28+
)
29+
)
30+
reveal_type(
31+
np.take( # E: Union[numpy.generic, datetime.datetime, datetime.timedelta, numpy.ndarray]
32+
B, [0]
33+
)
34+
)
35+
36+
reveal_type(np.reshape(a, 1)) # E: numpy.ndarray
37+
reveal_type(np.reshape(b, 1)) # E: numpy.ndarray
38+
reveal_type(np.reshape(c, 1)) # E: numpy.ndarray
39+
reveal_type(np.reshape(A, 1)) # E: numpy.ndarray
40+
reveal_type(np.reshape(B, 1)) # E: numpy.ndarray
41+
42+
reveal_type(np.choose(a, [True])) # E: numpy.bool_
43+
reveal_type(np.choose(b, [1.0])) # E: numpy.float32
44+
reveal_type(
45+
np.choose( # E: Union[numpy.generic, datetime.datetime, datetime.timedelta]
46+
c, [1.0]
47+
)
48+
)
49+
reveal_type(np.choose(A, [True])) # E: numpy.ndarray
50+
reveal_type(np.choose(B, [1.0])) # E: numpy.ndarray
51+
52+
reveal_type(np.repeat(a, 1)) # E: numpy.ndarray
53+
reveal_type(np.repeat(b, 1)) # E: numpy.ndarray
54+
reveal_type(np.repeat(c, 1)) # E: numpy.ndarray
55+
reveal_type(np.repeat(A, 1)) # E: numpy.ndarray
56+
reveal_type(np.repeat(B, 1)) # E: numpy.ndarray
57+
58+
# TODO: Add tests for np.put()
59+
60+
reveal_type(np.swapaxes(A, 0, 0)) # E: numpy.ndarray
61+
reveal_type(np.swapaxes(B, 0, 0)) # E: numpy.ndarray
62+
63+
reveal_type(np.transpose(a)) # E: numpy.ndarray
64+
reveal_type(np.transpose(b)) # E: numpy.ndarray
65+
reveal_type(np.transpose(c)) # E: numpy.ndarray
66+
reveal_type(np.transpose(A)) # E: numpy.ndarray
67+
reveal_type(np.transpose(B)) # E: numpy.ndarray
68+
69+
reveal_type(np.partition(a, 0)) # E: numpy.ndarray
70+
reveal_type(np.partition(b, 0)) # E: numpy.ndarray
71+
reveal_type(np.partition(c, 0)) # E: numpy.ndarray
72+
reveal_type(np.partition(A, 0)) # E: numpy.ndarray
73+
reveal_type(np.partition(B, 0)) # E: numpy.ndarray
74+
75+
reveal_type(np.argpartition(a, 0)) # E: numpy.ndarray
76+
reveal_type(np.argpartition(b, 0)) # E: numpy.ndarray
77+
reveal_type(np.argpartition(c, 0)) # E: numpy.ndarray
78+
reveal_type(np.argpartition(A, 0)) # E: numpy.ndarray
79+
reveal_type(np.argpartition(B, 0)) # E: numpy.ndarray
80+
81+
reveal_type(np.sort(A, 0)) # E: numpy.ndarray
82+
reveal_type(np.sort(B, 0)) # E: numpy.ndarray
83+
84+
reveal_type(np.argsort(A, 0)) # E: numpy.ndarray
85+
reveal_type(np.argsort(B, 0)) # E: numpy.ndarray

0 commit comments

Comments
 (0)