-
-
Notifications
You must be signed in to change notification settings - Fork 32
ENH: Add type annotations for the np.core.fromnumeric module: part 2/4 #71
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -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 | ||
|
||
|
@@ -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[ | ||
|
@@ -919,3 +920,69 @@ def argsort( | |
kind: Optional[_SortKind] = ..., | ||
order: Union[None, str, Sequence[str]] = ..., | ||
) -> ndarray: ... | ||
@overload | ||
def argmax( | ||
a: Union[Sequence[ArrayLike], ndarray], | ||
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 | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 But, I haven't complained about this on previous PRs, so maybe now is not the time to start objecting. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 |
||
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: ... |
There was a problem hiding this comment.
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?There was a problem hiding this comment.
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()
anddiagonal()
are one of the few cases where, at the very least, a vector or matrix is required. Hence the wholeSequence[ArrayLike]
andSequence[Sequence[ArrayLike]]
syntaxes.