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

ENH: Add type annotations for the np.core.fromnumeric module: part 2/4 #71

Merged
merged 3 commits into from
Jun 5, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 69 additions & 2 deletions numpy-stubs/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -772,6 +772,7 @@ _Mode = Literal["raise", "wrap", "clip"]
_Order = Literal["C", "F", "A"]
_PartitionKind = Literal["introselect"]
_SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"]
_Side = Literal["left", "right"]

# Various annotations for scalars

Expand All @@ -793,8 +794,8 @@ _ScalarGenericDT = TypeVar(
_Int = Union[int, integer]
_Bool = Union[bool, bool_]
_IntOrBool = Union[_Int, _Bool]
_ArrayLikeIntNested = Any # TODO: wait for support for recursive types
_ArrayLikeBoolNested = Any # TODO: wait for support for recursive types
_ArrayLikeIntNested = ArrayLike # TODO: wait for support for recursive types
_ArrayLikeBoolNested = ArrayLike # TODO: wait for support for recursive types

# Integers and booleans can generally be used interchangeably
_ArrayLikeIntOrBool = Union[
Expand Down Expand Up @@ -919,3 +920,69 @@ def argsort(
kind: Optional[_SortKind] = ...,
order: Union[None, str, Sequence[str]] = ...,
) -> ndarray: ...
@overload
def argmax(
a: Union[Sequence[ArrayLike], ndarray],
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not ArrayLike here and below?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So the likes of argmax(), argmin() and diagonal() are one of the few cases where, at the very least, a vector or matrix is required. Hence the whole Sequence[ArrayLike] and Sequence[Sequence[ArrayLike]] syntaxes.

axis: None = ...,
out: Optional[ndarray] = ...,
) -> integer: ...
@overload
def argmax(
a: Union[Sequence[ArrayLike], ndarray],
axis: int = ...,
out: Optional[ndarray] = ...,
) -> Union[integer, ndarray]: ...
@overload
def argmin(
a: Union[Sequence[ArrayLike], ndarray],
axis: None = ...,
out: Optional[ndarray] = ...,
) -> integer: ...
@overload
def argmin(
a: Union[Sequence[ArrayLike], ndarray],
axis: int = ...,
out: Optional[ndarray] = ...,
) -> Union[integer, ndarray]: ...
@overload
def searchsorted(
a: Union[Sequence[ArrayLike], ndarray],
v: _Scalar,
side: _Side = ...,
sorter: Union[None, Sequence[_IntOrBool], ndarray] = ..., # 1D int array
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't love the complexity things like this add to each signature versus having some uniform "array like but not a scalar" type that we can use everywhere. This is less detailed for now, but we can reach for this level of detail when making ndarray/ArrayLike generic; i.e. we could have something like _NonScalarArrayLike[T] and keep the code more readable/maintainable. Right now each signature ends up being its own little puzzle.

But, I haven't complained about this on previous PRs, so maybe now is not the time to start objecting.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I do like the _NonScalarArrayLike[T] suggestion.

A question though: do you feel it would be better to implement such a change in this pull request? Or shall i create a new pull request afterwards, addressing the issue from both #71 and #67?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let’s wait and I’ll merge this as-is. For now I’d like to prioritize getting stuff merged into NumPy, and then we can resume from there.

) -> integer: ...
@overload
def searchsorted(
a: Union[Sequence[ArrayLike], ndarray],
v: ArrayLike,
side: _Side = ...,
sorter: Union[None, Sequence[_IntOrBool], ndarray] = ..., # 1D int array
) -> ndarray: ...
def resize(a: ArrayLike, new_shape: _ShapeLike) -> ndarray: ...
@overload
def squeeze(a: _ScalarGeneric, axis: Optional[_ShapeLike] = ...) -> _ScalarGeneric: ...
@overload
def squeeze(a: ArrayLike, axis: Optional[_ShapeLike] = ...) -> ndarray: ...
def diagonal(
a: Union[Sequence[Sequence[ArrayLike]], ndarray], # >= 2D array
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we are encountering diminishing returns and increasing code complexity with things like Sequence[Sequence[ArrayLike]].

offset: int = ...,
axis1: int = ...,
axis2: int = ...,
) -> ndarray: ...
def trace(
a: Union[Sequence[Sequence[ArrayLike]], ndarray], # >= 2D array
offset: int = ...,
axis1: int = ...,
axis2: int = ...,
dtype: DtypeLike = ...,
out: Optional[ndarray] = ...,
) -> Union[number, ndarray]: ...
def ravel(a: ArrayLike, order: _Order = ...) -> ndarray: ...
def nonzero(a: ArrayLike) -> Tuple[ndarray, ...]: ...
def shape(a: ArrayLike) -> _Shape: ...
def compress(
condition: Union[Sequence[_Bool], ndarray], # 1D bool array
a: ArrayLike,
axis: Optional[int] = ...,
out: Optional[ndarray] = ...,
) -> ndarray: ...
37 changes: 37 additions & 0 deletions tests/fail/fromnumeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,3 +62,40 @@
np.argsort(A, axis="bob") # E: Argument "axis" to "argsort" has incompatible type
np.argsort(A, kind="bob") # E: Argument "kind" to "argsort" has incompatible type
np.argsort(A, order=range(5)) # E: Argument "order" to "argsort" has incompatible type

np.argmax(a) # E: No overload variant of "argmax" matches argument type
np.argmax(A, axis="bob") # E: No overload variant of "argmax" matches argument type
np.argmax(A, kind="bob") # E: No overload variant of "argmax" matches argument type

np.argmin(a) # E: No overload variant of "argmin" matches argument type
np.argmin(A, axis="bob") # E: No overload variant of "argmin" matches argument type
np.argmin(A, kind="bob") # E: No overload variant of "argmin" matches argument type

np.searchsorted(a, 0) # E: No overload variant of "searchsorted" matches argument type
np.searchsorted( # E: No overload variant of "searchsorted" matches argument type
A[0], 0, side="bob"
)
np.searchsorted( # E: No overload variant of "searchsorted" matches argument type
A[0], 0, sorter=1.0
)

np.resize(A, 1.0) # E: Argument 2 to "resize" has incompatible type

np.squeeze(A, 1.0) # E: No overload variant of "squeeze" matches argument type

np.diagonal(a) # E: Argument 1 to "diagonal" has incompatible type
np.diagonal(A, offset=None) # E: Argument "offset" to "diagonal" has incompatible type
np.diagonal(A, axis1="bob") # E: Argument "axis1" to "diagonal" has incompatible type
np.diagonal(A, axis2=[]) # E: Argument "axis2" to "diagonal" has incompatible type

np.trace(a) # E: Argument 1 to "trace" has incompatible type
np.trace(A, offset=None) # E: Argument "offset" to "trace" has incompatible type
np.trace(A, axis1="bob") # E: Argument "axis1" to "trace" has incompatible type
np.trace(A, axis2=[]) # E: Argument "axis2" to "trace" has incompatible type

np.ravel(a, order="bob") # E: Argument "order" to "ravel" has incompatible type

np.compress(True, A) # E: Argument 1 to "compress" has incompatible type
np.compress(
[True], A, axis=1.0 # E: Argument "axis" to "compress" has incompatible type
)
57 changes: 57 additions & 0 deletions tests/pass/fromnumeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,3 +60,60 @@

np.argsort(A, 0)
np.argsort(B, 0)

np.argmax(A)
np.argmax(B)
np.argmax(A, axis=0)
np.argmax(B, axis=0)

np.argmin(A)
np.argmin(B)
np.argmin(A, axis=0)
np.argmin(B, axis=0)

np.searchsorted(A[0], 0)
np.searchsorted(B[0], 0)
np.searchsorted(A[0], [0])
np.searchsorted(B[0], [0])

np.resize(a, (5, 5))
np.resize(b, (5, 5))
np.resize(c, (5, 5))
np.resize(A, (5, 5))
np.resize(B, (5, 5))

np.squeeze(a)
np.squeeze(b)
np.squeeze(c)
np.squeeze(A)
np.squeeze(B)

np.diagonal(A)
np.diagonal(B)

np.trace(A)
np.trace(B)

np.ravel(a)
np.ravel(b)
np.ravel(c)
np.ravel(A)
np.ravel(B)

np.nonzero(a)
np.nonzero(b)
np.nonzero(c)
np.nonzero(A)
np.nonzero(B)

np.shape(a)
np.shape(b)
np.shape(c)
np.shape(A)
np.shape(B)

np.compress([True], a)
np.compress([True], b)
np.compress([True], c)
np.compress([True], A)
np.compress([True], B)
57 changes: 57 additions & 0 deletions tests/reveal/fromnumeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,3 +76,60 @@

reveal_type(np.argsort(A, 0)) # E: numpy.ndarray
reveal_type(np.argsort(B, 0)) # E: numpy.ndarray

reveal_type(np.argmax(A)) # E: numpy.integer
reveal_type(np.argmax(B)) # E: numpy.integer
reveal_type(np.argmax(A, axis=0)) # E: Union[numpy.integer, numpy.ndarray]
reveal_type(np.argmax(B, axis=0)) # E: Union[numpy.integer, numpy.ndarray]

reveal_type(np.argmin(A)) # E: numpy.integer
reveal_type(np.argmin(B)) # E: numpy.integer
reveal_type(np.argmin(A, axis=0)) # E: Union[numpy.integer, numpy.ndarray]
reveal_type(np.argmin(B, axis=0)) # E: Union[numpy.integer, numpy.ndarray]

reveal_type(np.searchsorted(A[0], 0)) # E: numpy.integer
reveal_type(np.searchsorted(B[0], 0)) # E: numpy.integer
reveal_type(np.searchsorted(A[0], [0])) # E: numpy.ndarray
reveal_type(np.searchsorted(B[0], [0])) # E: numpy.ndarray

reveal_type(np.resize(a, (5, 5))) # E: numpy.ndarray
reveal_type(np.resize(b, (5, 5))) # E: numpy.ndarray
reveal_type(np.resize(c, (5, 5))) # E: numpy.ndarray
reveal_type(np.resize(A, (5, 5))) # E: numpy.ndarray
reveal_type(np.resize(B, (5, 5))) # E: numpy.ndarray

reveal_type(np.squeeze(a)) # E: numpy.bool_
reveal_type(np.squeeze(b)) # E: numpy.float32
reveal_type(np.squeeze(c)) # E: numpy.ndarray
reveal_type(np.squeeze(A)) # E: numpy.ndarray
reveal_type(np.squeeze(B)) # E: numpy.ndarray

reveal_type(np.diagonal(A)) # E: numpy.ndarray
reveal_type(np.diagonal(B)) # E: numpy.ndarray

reveal_type(np.trace(A)) # E: Union[numpy.number, numpy.ndarray]
reveal_type(np.trace(B)) # E: Union[numpy.number, numpy.ndarray]

reveal_type(np.ravel(a)) # E: numpy.ndarray
reveal_type(np.ravel(b)) # E: numpy.ndarray
reveal_type(np.ravel(c)) # E: numpy.ndarray
reveal_type(np.ravel(A)) # E: numpy.ndarray
reveal_type(np.ravel(B)) # E: numpy.ndarray

reveal_type(np.nonzero(a)) # E: tuple[numpy.ndarray]
reveal_type(np.nonzero(b)) # E: tuple[numpy.ndarray]
reveal_type(np.nonzero(c)) # E: tuple[numpy.ndarray]
reveal_type(np.nonzero(A)) # E: tuple[numpy.ndarray]
reveal_type(np.nonzero(B)) # E: tuple[numpy.ndarray]

reveal_type(np.shape(a)) # E: tuple[builtins.int]
reveal_type(np.shape(b)) # E: tuple[builtins.int]
reveal_type(np.shape(c)) # E: tuple[builtins.int]
reveal_type(np.shape(A)) # E: tuple[builtins.int]
reveal_type(np.shape(B)) # E: tuple[builtins.int]

reveal_type(np.compress([True], a)) # E: numpy.ndarray
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