diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/function_base.pyi | 184 | ||||
| -rw-r--r-- | numpy/typing/tests/data/fail/array_constructors.pyi | 6 | ||||
| -rw-r--r-- | numpy/typing/tests/data/reveal/array_constructors.pyi | 22 | ||||
| -rw-r--r-- | numpy/typing/tests/test_typing.py | 8 |
4 files changed, 187 insertions, 33 deletions
diff --git a/numpy/core/function_base.pyi b/numpy/core/function_base.pyi index 68d3b3a98..6e0843a0e 100644 --- a/numpy/core/function_base.pyi +++ b/numpy/core/function_base.pyi @@ -1,55 +1,195 @@ -from typing import overload, Tuple, Union, Sequence, Any, SupportsIndex, Literal, List +from typing import ( + Literal as L, + overload, + Tuple, + Union, + Any, + SupportsIndex, + List, + Type, + TypeVar, +) -from numpy import ndarray -from numpy.typing import ArrayLike, DTypeLike, _SupportsArray, _NumberLike_co +from numpy import floating, complexfloating, generic, dtype +from numpy.typing import ( + NDArray, + ArrayLike, + DTypeLike, + _SupportsDType, + _SupportsArray, + _NumberLike_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, +) -# TODO: wait for support for recursive types -_ArrayLikeNested = Sequence[Sequence[Any]] -_ArrayLikeNumber = Union[ - _NumberLike_co, Sequence[_NumberLike_co], ndarray, _SupportsArray, _ArrayLikeNested +_SCT = TypeVar("_SCT", bound=generic) + +_DTypeLike = Union[ + dtype[_SCT], + Type[_SCT], + _SupportsDType[dtype[_SCT]], ] __all__: List[str] @overload def linspace( - start: _ArrayLikeNumber, - stop: _ArrayLikeNumber, + start: _ArrayLikeFloat_co, + stop: _ArrayLikeFloat_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: L[False] = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> NDArray[floating[Any]]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: L[False] = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: L[False] = ..., + dtype: _DTypeLike[_SCT] = ..., + axis: SupportsIndex = ..., +) -> NDArray[_SCT]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, num: SupportsIndex = ..., endpoint: bool = ..., - retstep: Literal[False] = ..., + retstep: L[False] = ..., dtype: DTypeLike = ..., axis: SupportsIndex = ..., -) -> ndarray: ... +) -> NDArray[Any]: ... @overload def linspace( - start: _ArrayLikeNumber, - stop: _ArrayLikeNumber, + start: _ArrayLikeFloat_co, + stop: _ArrayLikeFloat_co, num: SupportsIndex = ..., endpoint: bool = ..., - retstep: Literal[True] = ..., + retstep: L[True] = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> Tuple[NDArray[floating[Any]], floating[Any]]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: L[True] = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> Tuple[NDArray[complexfloating[Any, Any]], complexfloating[Any, Any]]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: L[True] = ..., + dtype: _DTypeLike[_SCT] = ..., + axis: SupportsIndex = ..., +) -> Tuple[NDArray[_SCT], _SCT]: ... +@overload +def linspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + retstep: L[True] = ..., dtype: DTypeLike = ..., axis: SupportsIndex = ..., -) -> Tuple[ndarray, Any]: ... +) -> Tuple[NDArray[Any], Any]: ... +@overload def logspace( - start: _ArrayLikeNumber, - stop: _ArrayLikeNumber, + start: _ArrayLikeFloat_co, + stop: _ArrayLikeFloat_co, num: SupportsIndex = ..., endpoint: bool = ..., - base: _ArrayLikeNumber = ..., + base: _ArrayLikeFloat_co = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> NDArray[floating[Any]]: ... +@overload +def logspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + base: _ArrayLikeComplex_co = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def logspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + base: _ArrayLikeComplex_co = ..., + dtype: _DTypeLike[_SCT] = ..., + axis: SupportsIndex = ..., +) -> NDArray[_SCT]: ... +@overload +def logspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + base: _ArrayLikeComplex_co = ..., dtype: DTypeLike = ..., axis: SupportsIndex = ..., -) -> ndarray: ... +) -> NDArray[Any]: ... +@overload +def geomspace( + start: _ArrayLikeFloat_co, + stop: _ArrayLikeFloat_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> NDArray[floating[Any]]: ... +@overload +def geomspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + dtype: None = ..., + axis: SupportsIndex = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def geomspace( + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, + num: SupportsIndex = ..., + endpoint: bool = ..., + dtype: _DTypeLike[_SCT] = ..., + axis: SupportsIndex = ..., +) -> NDArray[_SCT]: ... +@overload def geomspace( - start: _ArrayLikeNumber, - stop: _ArrayLikeNumber, + start: _ArrayLikeComplex_co, + stop: _ArrayLikeComplex_co, num: SupportsIndex = ..., endpoint: bool = ..., dtype: DTypeLike = ..., axis: SupportsIndex = ..., -) -> ndarray: ... +) -> NDArray[Any]: ... # Re-exported to `np.lib.function_base` def add_newdoc( diff --git a/numpy/typing/tests/data/fail/array_constructors.pyi b/numpy/typing/tests/data/fail/array_constructors.pyi index 4f0a60b5b..065b7d8a0 100644 --- a/numpy/typing/tests/data/fail/array_constructors.pyi +++ b/numpy/typing/tests/data/fail/array_constructors.pyi @@ -21,10 +21,10 @@ np.linspace(0, 2, retstep=b'False') # E: No overload variant np.linspace(0, 2, dtype=0) # E: No overload variant np.linspace(0, 2, axis=None) # E: No overload variant -np.logspace(None, 'bob') # E: Argument 1 -np.logspace(0, 2, base=None) # E: Argument "base" +np.logspace(None, 'bob') # E: No overload variant +np.logspace(0, 2, base=None) # E: No overload variant -np.geomspace(None, 'bob') # E: Argument 1 +np.geomspace(None, 'bob') # E: No overload variant np.stack(generator) # E: No overload variant np.hstack({1, 2}) # E: No overload variant diff --git a/numpy/typing/tests/data/reveal/array_constructors.pyi b/numpy/typing/tests/data/reveal/array_constructors.pyi index 233988e63..ba5710e0f 100644 --- a/numpy/typing/tests/data/reveal/array_constructors.pyi +++ b/numpy/typing/tests/data/reveal/array_constructors.pyi @@ -114,10 +114,24 @@ reveal_type(np.require(B, requirements="W")) # E: SubClass[{float64}] reveal_type(np.require(B, requirements="A")) # E: SubClass[{float64}] reveal_type(np.require(C)) # E: ndarray[Any, Any] -reveal_type(np.linspace(0, 10)) # E: ndarray[Any, Any] -reveal_type(np.linspace(0, 10, retstep=True)) # E: Tuple[ndarray[Any, Any], Any] -reveal_type(np.logspace(0, 10)) # E: ndarray[Any, Any] -reveal_type(np.geomspace(1, 10)) # E: ndarray[Any, Any] +reveal_type(np.linspace(0, 10)) # E: ndarray[Any, dtype[floating[Any]]] +reveal_type(np.linspace(0, 10j)) # E: ndarray[Any, dtype[complexfloating[Any, Any]]] +reveal_type(np.linspace(0, 10, dtype=np.int64)) # E: ndarray[Any, dtype[{int64}]] +reveal_type(np.linspace(0, 10, dtype=int)) # E: ndarray[Any, dtype[Any]] +reveal_type(np.linspace(0, 10, retstep=True)) # E: Tuple[ndarray[Any, dtype[floating[Any]]], floating[Any]] +reveal_type(np.linspace(0j, 10, retstep=True)) # E: Tuple[ndarray[Any, dtype[complexfloating[Any, Any]]], complexfloating[Any, Any]] +reveal_type(np.linspace(0, 10, retstep=True, dtype=np.int64)) # E: Tuple[ndarray[Any, dtype[{int64}]], {int64}] +reveal_type(np.linspace(0j, 10, retstep=True, dtype=int)) # E: Tuple[ndarray[Any, dtype[Any]], Any] + +reveal_type(np.logspace(0, 10)) # E: ndarray[Any, dtype[floating[Any]]] +reveal_type(np.logspace(0, 10j)) # E: ndarray[Any, dtype[complexfloating[Any, Any]]] +reveal_type(np.logspace(0, 10, dtype=np.int64)) # E: ndarray[Any, dtype[{int64}]] +reveal_type(np.logspace(0, 10, dtype=int)) # E: ndarray[Any, dtype[Any]] + +reveal_type(np.geomspace(0, 10)) # E: ndarray[Any, dtype[floating[Any]]] +reveal_type(np.geomspace(0, 10j)) # E: ndarray[Any, dtype[complexfloating[Any, Any]]] +reveal_type(np.geomspace(0, 10, dtype=np.int64)) # E: ndarray[Any, dtype[{int64}]] +reveal_type(np.geomspace(0, 10, dtype=int)) # E: ndarray[Any, dtype[Any]] reveal_type(np.zeros_like(A)) # E: ndarray[Any, dtype[{float64}]] reveal_type(np.zeros_like(C)) # E: ndarray[Any, dtype[Any]] diff --git a/numpy/typing/tests/test_typing.py b/numpy/typing/tests/test_typing.py index fe58a8f4c..bb3914434 100644 --- a/numpy/typing/tests/test_typing.py +++ b/numpy/typing/tests/test_typing.py @@ -136,7 +136,7 @@ def test_fail(path: str) -> None: output_mypy = OUTPUT_MYPY assert path in output_mypy for error_line in output_mypy[path]: - error_line = _strip_filename(error_line) + error_line = _strip_filename(error_line).split("\n", 1)[0] match = re.match( r"(?P<lineno>\d+): (error|note): .+$", error_line, @@ -368,6 +368,7 @@ Expression: {} Expected reveal: {!r} Observed reveal: {!r} """ +_STRIP_PATTERN = re.compile(r"(\w+\.)+(\w+)") def _test_reveal( @@ -378,9 +379,8 @@ def _test_reveal( lineno: int, ) -> None: """Error-reporting helper function for `test_reveal`.""" - strip_pattern = re.compile(r"(\w+\.)+(\w+)") - stripped_reveal = strip_pattern.sub(strip_func, reveal) - stripped_expected_reveal = strip_pattern.sub(strip_func, expected_reveal) + stripped_reveal = _STRIP_PATTERN.sub(strip_func, reveal) + stripped_expected_reveal = _STRIP_PATTERN.sub(strip_func, expected_reveal) if stripped_reveal not in stripped_expected_reveal: raise AssertionError( _REVEAL_MSG.format(lineno, |