diff options
| -rw-r--r-- | numpy/lib/function_base.pyi | 201 | ||||
| -rw-r--r-- | numpy/typing/tests/data/fail/lib_function_base.pyi | 24 | ||||
| -rw-r--r-- | numpy/typing/tests/data/reveal/lib_function_base.pyi | 35 | ||||
| -rw-r--r-- | numpy/typing/tests/test_typing.py | 4 |
4 files changed, 249 insertions, 15 deletions
diff --git a/numpy/lib/function_base.pyi b/numpy/lib/function_base.pyi index cbbc87e65..ca0bd23d6 100644 --- a/numpy/lib/function_base.pyi +++ b/numpy/lib/function_base.pyi @@ -14,6 +14,7 @@ from typing import ( Protocol, SupportsIndex, Iterable, + SupportsInt, ) if sys.version_info >= (3, 10): @@ -27,6 +28,9 @@ from numpy import ( generic, floating, complexfloating, + float64, + complex128, + timedelta64, object_, _OrderKACF, ) @@ -42,7 +46,10 @@ from numpy.typing import ( _SupportsArray, _ArrayLikeComplex_co, _ArrayLikeFloat_co, + _ArrayLikeTD64_co, _ArrayLikeObject_co, + _FloatLike_co, + _ComplexLike_co, ) from numpy.core.function_base import ( @@ -261,8 +268,24 @@ def diff( append: ArrayLike = ..., ) -> NDArray[Any]: ... -# TODO -def interp(x, xp, fp, left=..., right=..., period=...): ... +@overload +def interp( + x: _ArrayLikeFloat_co, + xp: _ArrayLikeFloat_co, + fp: _ArrayLikeFloat_co, + left: None | _FloatLike_co = ..., + right: None | _FloatLike_co = ..., + period: None | _FloatLike_co = ..., +) -> NDArray[float64]: ... +@overload +def interp( + x: _ArrayLikeFloat_co, + xp: _ArrayLikeFloat_co, + fp: _ArrayLikeComplex_co, + left: None | _ComplexLike_co = ..., + right: None | _ComplexLike_co = ..., + period: None | _FloatLike_co = ..., +) -> NDArray[complex128]: ... @overload def angle(z: _ArrayLikeFloat_co, deg: bool = ...) -> floating[Any]: ... @@ -308,17 +331,169 @@ def disp( linefeed: bool = ..., ) -> None: ... -def cov(m, y=..., rowvar=..., bias=..., ddof=..., fweights=..., aweights=..., *, dtype=...): ... -def corrcoef(x, y=..., rowvar=..., bias = ..., ddof = ..., *, dtype=...): ... -def blackman(M): ... -def bartlett(M): ... -def hanning(M): ... -def hamming(M): ... -def i0(x): ... -def kaiser(M, beta): ... -def sinc(x): ... -def msort(a): ... -def median(a, axis=..., out=..., overwrite_input=..., keepdims=...): ... +@overload +def cov( + m: _ArrayLikeFloat_co, + y: None | _ArrayLikeFloat_co = ..., + rowvar: bool = ..., + bias: bool = ..., + ddof: None | SupportsIndex | SupportsInt = ..., + fweights: None | ArrayLike = ..., + aweights: None | ArrayLike = ..., + *, + dtype: None = ..., +) -> NDArray[floating[Any]]: ... +@overload +def cov( + m: _ArrayLikeComplex_co, + y: None | _ArrayLikeComplex_co = ..., + rowvar: bool = ..., + bias: bool = ..., + ddof: None | SupportsIndex | SupportsInt = ..., + fweights: None | ArrayLike = ..., + aweights: None | ArrayLike = ..., + *, + dtype: None = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def cov( + m: _ArrayLikeComplex_co, + y: None | _ArrayLikeComplex_co = ..., + rowvar: bool = ..., + bias: bool = ..., + ddof: None | SupportsIndex | SupportsInt = ..., + fweights: None | ArrayLike = ..., + aweights: None | ArrayLike = ..., + *, + dtype: _DTypeLike[_SCT], +) -> NDArray[_SCT]: ... +@overload +def cov( + m: _ArrayLikeComplex_co, + y: None | _ArrayLikeComplex_co = ..., + rowvar: bool = ..., + bias: bool = ..., + ddof: None | SupportsIndex | SupportsInt = ..., + fweights: None | ArrayLike = ..., + aweights: None | ArrayLike = ..., + *, + dtype: DTypeLike, +) -> NDArray[Any]: ... + +# NOTE `bias` and `ddof` have been deprecated +@overload +def corrcoef( + m: _ArrayLikeFloat_co, + y: None | _ArrayLikeFloat_co = ..., + rowvar: bool = ..., + *, + dtype: None = ..., +) -> NDArray[floating[Any]]: ... +@overload +def corrcoef( + m: _ArrayLikeComplex_co, + y: None | _ArrayLikeComplex_co = ..., + rowvar: bool = ..., + *, + dtype: None = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def corrcoef( + m: _ArrayLikeComplex_co, + y: None | _ArrayLikeComplex_co = ..., + rowvar: bool = ..., + *, + dtype: _DTypeLike[_SCT], +) -> NDArray[_SCT]: ... +@overload +def corrcoef( + m: _ArrayLikeComplex_co, + y: None | _ArrayLikeComplex_co = ..., + rowvar: bool = ..., + *, + dtype: DTypeLike, +) -> NDArray[Any]: ... + +def blackman(M: _FloatLike_co) -> NDArray[floating[Any]]: ... + +def bartlett(M: _FloatLike_co) -> NDArray[floating[Any]]: ... + +def hanning(M: _FloatLike_co) -> NDArray[floating[Any]]: ... + +def hamming(M: _FloatLike_co) -> NDArray[floating[Any]]: ... + +def i0(x: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... + +def kaiser( + M: _FloatLike_co, + beta: _FloatLike_co, +) -> NDArray[floating[Any]]: ... + +@overload +def sinc(x: _FloatLike_co) -> floating[Any]: ... +@overload +def sinc(x: _ComplexLike_co) -> complexfloating[Any, Any]: ... +@overload +def sinc(x: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... +@overload +def sinc(x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + +@overload +def msort(a: _ArrayType) -> _ArrayType: ... +@overload +def msort(a: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... +@overload +def msort(a: ArrayLike) -> NDArray[Any]: ... + +@overload +def median( + a: _ArrayLikeFloat_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + keepdims: L[False] = ..., +) -> floating[Any]: ... +@overload +def median( + a: _ArrayLikeComplex_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + keepdims: L[False] = ..., +) -> complexfloating[Any, Any]: ... +@overload +def median( + a: _ArrayLikeTD64_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + keepdims: L[False] = ..., +) -> timedelta64: ... +@overload +def median( + a: _ArrayLikeObject_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + keepdims: L[False] = ..., +) -> Any: ... +@overload +def median( + a: _ArrayLikeFloat_co | _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, + axis: None | _ShapeLike = ..., + out: None = ..., + overwrite_input: bool = ..., + keepdims: bool = ..., +) -> Any: ... +@overload +def median( + a: _ArrayLikeFloat_co | _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, + axis: None | _ShapeLike = ..., + out: _ArrayType = ..., + overwrite_input: bool = ..., + keepdims: bool = ..., +) -> _ArrayType: ... + def percentile(a, q, axis=..., out=..., overwrite_input=..., interpolation=..., keepdims=...): ... def quantile(a, q, axis=..., out=..., overwrite_input=..., interpolation=..., keepdims=...): ... def trapz(y, x=..., dx=..., axis=...): ... diff --git a/numpy/typing/tests/data/fail/lib_function_base.pyi b/numpy/typing/tests/data/fail/lib_function_base.pyi index 019bd7f01..2a367b9fe 100644 --- a/numpy/typing/tests/data/fail/lib_function_base.pyi +++ b/numpy/typing/tests/data/fail/lib_function_base.pyi @@ -3,9 +3,11 @@ from typing import Any import numpy as np import numpy.typing as npt -AR_m: npt.NDArray[np.timedelta64] AR_f8: npt.NDArray[np.float64] AR_c16: npt.NDArray[np.complex128] +AR_m: npt.NDArray[np.timedelta64] +AR_M: npt.NDArray[np.datetime64] +AR_O: npt.NDArray[np.object_] np.average(AR_m) # E: incompatible type np.select(1, [AR_f8]) # E: incompatible type @@ -17,3 +19,23 @@ np.place(1, [True], 1.5) # E: incompatible type np.vectorize(1) # E: incompatible type np.add_newdoc("__main__", 1.5, "docstring") # E: incompatible type np.place(AR_f8, slice(None), 5) # E: incompatible type + +np.interp(AR_f8, AR_c16, AR_f8) # E: incompatible type +np.interp(AR_c16, AR_f8, AR_f8) # E: incompatible type +np.interp(AR_f8, AR_f8, AR_f8, period=AR_c16) # E: No overload variant +np.interp(AR_f8, AR_f8, AR_O) # E: incompatible type + +np.cov(AR_m) # E: incompatible type +np.cov(AR_O) # E: incompatible type +np.corrcoef(AR_m) # E: incompatible type +np.corrcoef(AR_O) # E: incompatible type +np.corrcoef(AR_f8, bias=True) # E: No overload variant +np.corrcoef(AR_f8, ddof=2) # E: No overload variant +np.blackman(1j) # E: incompatible type +np.bartlett(1j) # E: incompatible type +np.hanning(1j) # E: incompatible type +np.hamming(1j) # E: incompatible type +np.hamming(AR_c16) # E: incompatible type +np.kaiser(1j, 1) # E: incompatible type +np.sinc(AR_O) # E: incompatible type +np.median(AR_M) # E: incompatible type diff --git a/numpy/typing/tests/data/reveal/lib_function_base.pyi b/numpy/typing/tests/data/reveal/lib_function_base.pyi index 76d54c49f..81e6d52b2 100644 --- a/numpy/typing/tests/data/reveal/lib_function_base.pyi +++ b/numpy/typing/tests/data/reveal/lib_function_base.pyi @@ -11,6 +11,7 @@ AR_LIKE_f8: list[float] AR_i8: npt.NDArray[np.int64] AR_f8: npt.NDArray[np.float64] AR_c16: npt.NDArray[np.complex128] +AR_m: npt.NDArray[np.timedelta64] AR_O: npt.NDArray[np.object_] AR_b: npt.NDArray[np.bool_] AR_U: npt.NDArray[np.str_] @@ -97,3 +98,37 @@ reveal_type(np.place(AR_f8, mask=AR_i8, vals=5.0)) # E: None reveal_type(np.disp(1, linefeed=True)) # E: None with open("test", "w") as f: reveal_type(np.disp("message", device=f)) # E: None + +reveal_type(np.cov(AR_f8, bias=True)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] +reveal_type(np.cov(AR_f8, AR_c16, ddof=1)) # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]] +reveal_type(np.cov(AR_f8, aweights=AR_f8, dtype=np.float32)) # E: numpy.ndarray[Any, numpy.dtype[{float32}]] +reveal_type(np.cov(AR_f8, fweights=AR_f8, dtype=float)) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.corrcoef(AR_f8, rowvar=True)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] +reveal_type(np.corrcoef(AR_f8, AR_c16)) # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]] +reveal_type(np.corrcoef(AR_f8, dtype=np.float32)) # E: numpy.ndarray[Any, numpy.dtype[{float32}]] +reveal_type(np.corrcoef(AR_f8, dtype=float)) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.blackman(5)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] +reveal_type(np.bartlett(6)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] +reveal_type(np.hanning(4.5)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] +reveal_type(np.hamming(0)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] +reveal_type(np.i0(AR_i8)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] +reveal_type(np.kaiser(4, 5.9)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] + +reveal_type(np.sinc(1.0)) # E: numpy.floating[Any] +reveal_type(np.sinc(1j)) # E: numpy.complexfloating[Any, Any] +reveal_type(np.sinc(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] +reveal_type(np.sinc(AR_c16)) # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]] + +reveal_type(np.msort(CHAR_AR_U)) # E: Any +reveal_type(np.msort(AR_U)) # E: numpy.ndarray[Any, numpy.dtype[numpy.str_]] +reveal_type(np.msort(AR_LIKE_f8)) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.median(AR_f8, keepdims=False)) # E: numpy.floating[Any] +reveal_type(np.median(AR_c16, overwrite_input=True)) # E: numpy.complexfloating[Any, Any] +reveal_type(np.median(AR_m)) # E: numpy.timedelta64 +reveal_type(np.median(AR_O)) # E: Any +reveal_type(np.median(AR_f8, keepdims=True)) # E: Any +reveal_type(np.median(AR_c16, axis=0)) # E: Any +reveal_type(np.median(AR_LIKE_f8, out=AR_c16)) # E: numpy.ndarray[Any, numpy.dtype[{complex128}]] diff --git a/numpy/typing/tests/test_typing.py b/numpy/typing/tests/test_typing.py index 0f3e10b7b..f303ebea3 100644 --- a/numpy/typing/tests/test_typing.py +++ b/numpy/typing/tests/test_typing.py @@ -160,7 +160,9 @@ def test_fail(path: str) -> None: expected_error = errors.get(lineno) _test_fail(path, marker, expected_error, lineno) else: - pytest.fail(f"Unexpected mypy output\n\n{errors[lineno]}") + pytest.fail( + f"Unexpected mypy output at line {lineno}\n\n{errors[lineno]}" + ) _FAIL_MSG1 = """Extra error at line {} |