123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357 |
- import sys
- import datetime as dt
- from typing import Optional, Union, Sequence, Tuple, Any, overload, TypeVar
- from numpy import (
- ndarray,
- signedinteger,
- bool_,
- generic,
- _OrderKACF,
- _OrderACF,
- _ArrayLikeBool,
- _ArrayLikeIntOrBool,
- _ModeKind,
- _PartitionKind,
- _SortKind,
- _SortSide,
- )
- from numpy.typing import (
- DTypeLike,
- ArrayLike,
- _ShapeLike,
- _Shape,
- _NumberLike,
- )
- if sys.version_info >= (3, 8):
- from typing import Literal
- else:
- from typing_extensions import Literal
- # Various annotations for scalars
- # While dt.datetime and dt.timedelta are not technically part of NumPy,
- # they are one of the rare few builtin scalars which serve as valid return types.
- # See https://github.com/numpy/numpy-stubs/pull/67#discussion_r412604113.
- _ScalarNumpy = Union[generic, dt.datetime, dt.timedelta]
- _ScalarBuiltin = Union[str, bytes, dt.date, dt.timedelta, bool, int, float, complex]
- _Scalar = Union[_ScalarBuiltin, _ScalarNumpy]
- # Integers and booleans can generally be used interchangeably
- _ScalarGeneric = TypeVar("_ScalarGeneric", bound=generic)
- # The signature of take() follows a common theme with its overloads:
- # 1. A generic comes in; the same generic comes out
- # 2. A scalar comes in; a generic comes out
- # 3. An array-like object comes in; some keyword ensures that a generic comes out
- # 4. An array-like object comes in; an ndarray or generic comes out
- def take(
- a: ArrayLike,
- indices: _ArrayLikeIntOrBool,
- axis: Optional[int] = ...,
- out: Optional[ndarray] = ...,
- mode: _ModeKind = ...,
- ) -> Any: ...
- def reshape(
- a: ArrayLike,
- newshape: _ShapeLike,
- order: _OrderACF = ...,
- ) -> ndarray: ...
- def choose(
- a: _ArrayLikeIntOrBool,
- choices: ArrayLike,
- out: Optional[ndarray] = ...,
- mode: _ModeKind = ...,
- ) -> Any: ...
- def repeat(
- a: ArrayLike,
- repeats: _ArrayLikeIntOrBool,
- axis: Optional[int] = ...,
- ) -> ndarray: ...
- def put(
- a: ndarray,
- ind: _ArrayLikeIntOrBool,
- v: ArrayLike,
- mode: _ModeKind = ...,
- ) -> None: ...
- def swapaxes(
- a: ArrayLike,
- axis1: int,
- axis2: int,
- ) -> ndarray: ...
- def transpose(
- a: ArrayLike,
- axes: Union[None, Sequence[int], ndarray] = ...
- ) -> ndarray: ...
- def partition(
- a: ArrayLike,
- kth: _ArrayLikeIntOrBool,
- axis: Optional[int] = ...,
- kind: _PartitionKind = ...,
- order: Union[None, str, Sequence[str]] = ...,
- ) -> ndarray: ...
- def argpartition(
- a: ArrayLike,
- kth: _ArrayLikeIntOrBool,
- axis: Optional[int] = ...,
- kind: _PartitionKind = ...,
- order: Union[None, str, Sequence[str]] = ...,
- ) -> Any: ...
- def sort(
- a: ArrayLike,
- axis: Optional[int] = ...,
- kind: Optional[_SortKind] = ...,
- order: Union[None, str, Sequence[str]] = ...,
- ) -> ndarray: ...
- def argsort(
- a: ArrayLike,
- axis: Optional[int] = ...,
- kind: Optional[_SortKind] = ...,
- order: Union[None, str, Sequence[str]] = ...,
- ) -> ndarray: ...
- @overload
- def argmax(
- a: ArrayLike,
- axis: None = ...,
- out: Optional[ndarray] = ...,
- ) -> signedinteger[Any]: ...
- @overload
- def argmax(
- a: ArrayLike,
- axis: Optional[int] = ...,
- out: Optional[ndarray] = ...,
- ) -> Any: ...
- @overload
- def argmin(
- a: ArrayLike,
- axis: None = ...,
- out: Optional[ndarray] = ...,
- ) -> signedinteger[Any]: ...
- @overload
- def argmin(
- a: ArrayLike,
- axis: Optional[int] = ...,
- out: Optional[ndarray] = ...,
- ) -> Any: ...
- @overload
- def searchsorted(
- a: ArrayLike,
- v: _Scalar,
- side: _SortSide = ...,
- sorter: Optional[_ArrayLikeIntOrBool] = ..., # 1D int array
- ) -> signedinteger[Any]: ...
- @overload
- def searchsorted(
- a: ArrayLike,
- v: ArrayLike,
- side: _SortSide = ...,
- sorter: Optional[_ArrayLikeIntOrBool] = ..., # 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: ArrayLike,
- offset: int = ...,
- axis1: int = ...,
- axis2: int = ..., # >= 2D array
- ) -> ndarray: ...
- def trace(
- a: ArrayLike, # >= 2D array
- offset: int = ...,
- axis1: int = ...,
- axis2: int = ...,
- dtype: DTypeLike = ...,
- out: Optional[ndarray] = ...,
- ) -> Any: ...
- def ravel(a: ArrayLike, order: _OrderKACF = ...) -> ndarray: ...
- def nonzero(a: ArrayLike) -> Tuple[ndarray, ...]: ...
- def shape(a: ArrayLike) -> _Shape: ...
- def compress(
- condition: ArrayLike, # 1D bool array
- a: ArrayLike,
- axis: Optional[int] = ...,
- out: Optional[ndarray] = ...,
- ) -> ndarray: ...
- @overload
- def clip(
- a: ArrayLike,
- a_min: ArrayLike,
- a_max: Optional[ArrayLike],
- out: Optional[ndarray] = ...,
- **kwargs: Any,
- ) -> Any: ...
- @overload
- def clip(
- a: ArrayLike,
- a_min: None,
- a_max: ArrayLike,
- out: Optional[ndarray] = ...,
- **kwargs: Any,
- ) -> Any: ...
- def sum(
- a: ArrayLike,
- axis: _ShapeLike = ...,
- dtype: DTypeLike = ...,
- out: Optional[ndarray] = ...,
- keepdims: bool = ...,
- initial: _NumberLike = ...,
- where: _ArrayLikeBool = ...,
- ) -> Any: ...
- @overload
- def all(
- a: ArrayLike,
- axis: None = ...,
- out: None = ...,
- keepdims: Literal[False] = ...,
- ) -> bool_: ...
- @overload
- def all(
- a: ArrayLike,
- axis: Optional[_ShapeLike] = ...,
- out: Optional[ndarray] = ...,
- keepdims: bool = ...,
- ) -> Any: ...
- @overload
- def any(
- a: ArrayLike,
- axis: None = ...,
- out: None = ...,
- keepdims: Literal[False] = ...,
- ) -> bool_: ...
- @overload
- def any(
- a: ArrayLike,
- axis: Optional[_ShapeLike] = ...,
- out: Optional[ndarray] = ...,
- keepdims: bool = ...,
- ) -> Any: ...
- def cumsum(
- a: ArrayLike,
- axis: Optional[int] = ...,
- dtype: DTypeLike = ...,
- out: Optional[ndarray] = ...,
- ) -> ndarray: ...
- def ptp(
- a: ArrayLike,
- axis: Optional[_ShapeLike] = ...,
- out: Optional[ndarray] = ...,
- keepdims: bool = ...,
- ) -> Any: ...
- def amax(
- a: ArrayLike,
- axis: Optional[_ShapeLike] = ...,
- out: Optional[ndarray] = ...,
- keepdims: bool = ...,
- initial: _NumberLike = ...,
- where: _ArrayLikeBool = ...,
- ) -> Any: ...
- def amin(
- a: ArrayLike,
- axis: Optional[_ShapeLike] = ...,
- out: Optional[ndarray] = ...,
- keepdims: bool = ...,
- initial: _NumberLike = ...,
- where: _ArrayLikeBool = ...,
- ) -> Any: ...
- # TODO: `np.prod()``: For object arrays `initial` does not necessarily
- # have to be a numerical scalar.
- # The only requirement is that it is compatible
- # with the `.__mul__()` method(s) of the passed array's elements.
- # Note that the same situation holds for all wrappers around
- # `np.ufunc.reduce`, e.g. `np.sum()` (`.__add__()`).
- def prod(
- a: ArrayLike,
- axis: Optional[_ShapeLike] = ...,
- dtype: DTypeLike = ...,
- out: Optional[ndarray] = ...,
- keepdims: bool = ...,
- initial: _NumberLike = ...,
- where: _ArrayLikeBool = ...,
- ) -> Any: ...
- def cumprod(
- a: ArrayLike,
- axis: Optional[int] = ...,
- dtype: DTypeLike = ...,
- out: Optional[ndarray] = ...,
- ) -> ndarray: ...
- def ndim(a: ArrayLike) -> int: ...
- def size(a: ArrayLike, axis: Optional[int] = ...) -> int: ...
- def around(
- a: ArrayLike,
- decimals: int = ...,
- out: Optional[ndarray] = ...,
- ) -> Any: ...
- def mean(
- a: ArrayLike,
- axis: Optional[_ShapeLike] = ...,
- dtype: DTypeLike = ...,
- out: Optional[ndarray] = ...,
- keepdims: bool = ...,
- ) -> Any: ...
- def std(
- a: ArrayLike,
- axis: Optional[_ShapeLike] = ...,
- dtype: DTypeLike = ...,
- out: Optional[ndarray] = ...,
- ddof: int = ...,
- keepdims: bool = ...,
- ) -> Any: ...
- def var(
- a: ArrayLike,
- axis: Optional[_ShapeLike] = ...,
- dtype: DTypeLike = ...,
- out: Optional[ndarray] = ...,
- ddof: int = ...,
- keepdims: bool = ...,
- ) -> Any: ...
|