From dcb39277c201a418c653a1a98f9a6680fb3d5495 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Sun, 23 Aug 2020 15:36:07 +0200 Subject: ENH: Add stubs for `np.core.function_base` --- numpy/core/function_base.pyi | 53 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 numpy/core/function_base.pyi (limited to 'numpy/core/function_base.pyi') diff --git a/numpy/core/function_base.pyi b/numpy/core/function_base.pyi new file mode 100644 index 000000000..2edcb88c3 --- /dev/null +++ b/numpy/core/function_base.pyi @@ -0,0 +1,53 @@ +import sys +from typing import overload, Tuple, Union, Sequence, Any + +from numpy import ndarray, floating, number, _NumberLike +from numpy.typing import ArrayLike, DtypeLike, _SupportsArray + +if sys.version_info >= (3, 8): + from typing import SupportsIndex, Literal +else: + from typing_extensions import SupportsIndex, Literal + +# TODO: wait for support for recursive types +_ArrayLikeNested = Sequence[Sequence[Any]] +_ArrayLikeNumber = Union[ + _NumberLike, Sequence[_NumberLike], ndarray, _SupportsArray, _ArrayLikeNested +] +@overload +def linspace( + start: _ArrayLikeNumber, + stop: _ArrayLikeNumber, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: Literal[False] = ..., + dtype: DtypeLike = ..., + axis: int = ..., +) -> ndarray: ... +@overload +def linspace( + start: _ArrayLikeNumber, + stop: _ArrayLikeNumber, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: Literal[True] = ..., + dtype: DtypeLike = ..., + axis: int = ..., +) -> Tuple[ndarray, floating]: ... +def logspace( + start: _ArrayLikeNumber, + stop: _ArrayLikeNumber, + num: SupportsIndex = ..., + endpoint: bool = ..., + base: _ArrayLikeNumber = ..., + dtype: DtypeLike = ..., + axis: int = ..., +) -> ndarray: ... +def geomspace( + start: _ArrayLikeNumber, + stop: _ArrayLikeNumber, + num: SupportsIndex = ..., + endpoint: bool = ..., + dtype: DtypeLike = ..., + axis: int = ..., +) -> ndarray: ... -- cgit v1.2.1