diff options
| -rw-r--r-- | numpy/lib/twodim_base.pyi | 265 | ||||
| -rw-r--r-- | numpy/typing/tests/data/fail/twodim_base.py | 37 | ||||
| -rw-r--r-- | numpy/typing/tests/data/reveal/twodim_base.py | 72 |
3 files changed, 353 insertions, 21 deletions
diff --git a/numpy/lib/twodim_base.pyi b/numpy/lib/twodim_base.pyi index 79b9511b8..007338d77 100644 --- a/numpy/lib/twodim_base.pyi +++ b/numpy/lib/twodim_base.pyi @@ -1,32 +1,255 @@ -from typing import List, Optional, Any +from typing import ( + Any, + Callable, + List, + Sequence, + overload, + Tuple, + Type, + TypeVar, + Union, +) -from numpy import ndarray, _OrderCF -from numpy.typing import ArrayLike, DTypeLike +from numpy import ( + ndarray, + dtype, + generic, + number, + bool_, + timedelta64, + datetime64, + int_, + intp, + float64, + signedinteger, + floating, + complexfloating, + object_, + _OrderCF, +) + +from numpy.typing import ( + DTypeLike, + _SupportsDType, + ArrayLike, + NDArray, + _NestedSequence, + _SupportsArray, + _ArrayLikeInt_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeObject_co, +) + +_T = TypeVar("_T") +_SCT = TypeVar("_SCT", bound=generic) + +# The returned arrays dtype must be compatible with `np.equal` +_MaskFunc = Callable[ + [NDArray[int_], _T], + NDArray[Union[number[Any], bool_, timedelta64, datetime64, object_]], +] + +_DTypeLike = Union[ + Type[_SCT], + dtype[_SCT], + _SupportsDType[dtype[_SCT]], +] +_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]] __all__: List[str] -def fliplr(m): ... -def flipud(m): ... +@overload +def fliplr(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... +@overload +def fliplr(m: ArrayLike) -> NDArray[Any]: ... + +@overload +def flipud(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... +@overload +def flipud(m: ArrayLike) -> NDArray[Any]: ... +@overload def eye( N: int, - M: Optional[int] = ..., + M: None | int = ..., + k: int = ..., + dtype: None = ..., + order: _OrderCF = ..., + *, + like: None | ArrayLike = ..., +) -> NDArray[float64]: ... +@overload +def eye( + N: int, + M: None | int = ..., + k: int = ..., + dtype: _DTypeLike[_SCT] = ..., + order: _OrderCF = ..., + *, + like: None | ArrayLike = ..., +) -> NDArray[_SCT]: ... +@overload +def eye( + N: int, + M: None | int = ..., k: int = ..., dtype: DTypeLike = ..., order: _OrderCF = ..., *, - like: Optional[ArrayLike] = ... -) -> ndarray[Any, Any]: ... - -def diag(v, k=...): ... -def diagflat(v, k=...): ... -def tri(N, M=..., k=..., dtype = ..., *, like=...): ... -def tril(m, k=...): ... -def triu(m, k=...): ... -def vander(x, N=..., increasing=...): ... -def histogram2d(x, y, bins=..., range=..., normed=..., weights=..., density=...): ... -def mask_indices(n, mask_func, k=...): ... -def tril_indices(n, k=..., m=...): ... -def tril_indices_from(arr, k=...): ... -def triu_indices(n, k=..., m=...): ... -def triu_indices_from(arr, k=...): ... + like: None | ArrayLike = ..., +) -> NDArray[Any]: ... + +@overload +def diag(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... +@overload +def diag(v: ArrayLike, k: int = ...) -> NDArray[Any]: ... + +@overload +def diagflat(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... +@overload +def diagflat(v: ArrayLike, k: int = ...) -> NDArray[Any]: ... + +@overload +def tri( + N: int, + M: None | int = ..., + k: int = ..., + dtype: None = ..., + *, + like: None | ArrayLike = ... +) -> NDArray[float64]: ... +@overload +def tri( + N: int, + M: None | int = ..., + k: int = ..., + dtype: _DTypeLike[_SCT] = ..., + *, + like: None | ArrayLike = ... +) -> NDArray[_SCT]: ... +@overload +def tri( + N: int, + M: None | int = ..., + k: int = ..., + dtype: DTypeLike = ..., + *, + like: None | ArrayLike = ... +) -> NDArray[Any]: ... + +@overload +def tril(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... +@overload +def tril(v: ArrayLike, k: int = ...) -> NDArray[Any]: ... + +@overload +def triu(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... +@overload +def triu(v: ArrayLike, k: int = ...) -> NDArray[Any]: ... + +@overload +def vander( # type: ignore[misc] + x: _ArrayLikeInt_co, + N: None | int = ..., + increasing: bool = ..., +) -> NDArray[signedinteger[Any]]: ... +@overload +def vander( # type: ignore[misc] + x: _ArrayLikeFloat_co, + N: None | int = ..., + increasing: bool = ..., +) -> NDArray[floating[Any]]: ... +@overload +def vander( + x: _ArrayLikeComplex_co, + N: None | int = ..., + increasing: bool = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def vander( + x: _ArrayLikeObject_co, + N: None | int = ..., + increasing: bool = ..., +) -> NDArray[object_]: ... + +@overload +def histogram2d( # type: ignore[misc] + x: _ArrayLikeFloat_co, + y: _ArrayLikeFloat_co, + bins: int | Sequence[int] = ..., + range: None | _ArrayLikeFloat_co = ..., + normed: None | bool = ..., + weights: None | _ArrayLikeFloat_co = ..., + density: None | bool = ..., +) -> Tuple[ + NDArray[float64], + NDArray[floating[Any]], + NDArray[floating[Any]], +]: ... +@overload +def histogram2d( + x: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co, + bins: int | Sequence[int] = ..., + range: None | _ArrayLikeFloat_co = ..., + normed: None | bool = ..., + weights: None | _ArrayLikeFloat_co = ..., + density: None | bool = ..., +) -> Tuple[ + NDArray[float64], + NDArray[complexfloating[Any, Any]], + NDArray[complexfloating[Any, Any]], +]: ... +@overload # TODO: Sort out `bins` +def histogram2d( + x: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co, + bins: Sequence[_ArrayLikeInt_co], + range: None | _ArrayLikeFloat_co = ..., + normed: None | bool = ..., + weights: None | _ArrayLikeFloat_co = ..., + density: None | bool = ..., +) -> Tuple[ + NDArray[float64], + NDArray[Any], + NDArray[Any], +]: ... + +# NOTE: we're assuming/demanding here the `mask_func` returns +# an ndarray of shape `(n, n)`; otherwise there is the possibility +# of the output tuple having more or less than 2 elements +@overload +def mask_indices( + n: int, + mask_func: _MaskFunc[int], + k: int = ..., +) -> Tuple[NDArray[intp], NDArray[intp]]: ... +@overload +def mask_indices( + n: int, + mask_func: _MaskFunc[_T], + k: _T, +) -> Tuple[NDArray[intp], NDArray[intp]]: ... + +def tril_indices( + n: int, + k: int = ..., + m: None | int = ..., +) -> Tuple[NDArray[int_], NDArray[int_]]: ... + +def tril_indices_from( + arr: NDArray[Any], + k: int = ..., +) -> Tuple[NDArray[int_], NDArray[int_]]: ... + +def triu_indices( + n: int, + k: int = ..., + m: None | int = ..., +) -> Tuple[NDArray[int_], NDArray[int_]]: ... + +def triu_indices_from( + arr: NDArray[Any], + k: int = ..., +) -> Tuple[NDArray[int_], NDArray[int_]]: ... diff --git a/numpy/typing/tests/data/fail/twodim_base.py b/numpy/typing/tests/data/fail/twodim_base.py new file mode 100644 index 000000000..ab34a374c --- /dev/null +++ b/numpy/typing/tests/data/fail/twodim_base.py @@ -0,0 +1,37 @@ +from typing import Any, List, TypeVar + +import numpy as np +import numpy.typing as npt + + +def func1(ar: npt.NDArray[Any], a: int) -> npt.NDArray[np.str_]: + pass + + +def func2(ar: npt.NDArray[Any], a: float) -> float: + pass + + +AR_b: npt.NDArray[np.bool_] +AR_m: npt.NDArray[np.timedelta64] + +AR_LIKE_b: List[bool] + +np.eye(10, M=20.0) # E: No overload variant +np.eye(10, k=2.5, dtype=int) # E: No overload variant + +np.diag(AR_b, k=0.5) # E: No overload variant +np.diagflat(AR_b, k=0.5) # E: No overload variant + +np.tri(10, M=20.0) # E: No overload variant +np.tri(10, k=2.5, dtype=int) # E: No overload variant + +np.tril(AR_b, k=0.5) # E: No overload variant +np.triu(AR_b, k=0.5) # E: No overload variant + +np.vander(AR_m) # E: incompatible type + +np.histogram2d(AR_m) # E: No overload variant + +np.mask_indices(10, func1) # E: incompatible type +np.mask_indices(10, func2, 10.5) # E: incompatible type diff --git a/numpy/typing/tests/data/reveal/twodim_base.py b/numpy/typing/tests/data/reveal/twodim_base.py new file mode 100644 index 000000000..b95fbc71e --- /dev/null +++ b/numpy/typing/tests/data/reveal/twodim_base.py @@ -0,0 +1,72 @@ +from typing import Any, List, TypeVar + +import numpy as np +import numpy.typing as npt + +_SCT = TypeVar("_SCT", bound=np.generic) + + +def func1(ar: npt.NDArray[_SCT], a: int) -> npt.NDArray[_SCT]: + pass + + +def func2(ar: npt.NDArray[np.number[Any]], a: str) -> npt.NDArray[np.float64]: + pass + + +AR_b: npt.NDArray[np.bool_] +AR_u: npt.NDArray[np.uint64] +AR_i: npt.NDArray[np.int64] +AR_f: npt.NDArray[np.float64] +AR_c: npt.NDArray[np.complex128] +AR_O: npt.NDArray[np.object_] + +AR_LIKE_b: List[bool] + +reveal_type(np.fliplr(AR_b)) # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]] +reveal_type(np.fliplr(AR_LIKE_b)) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.flipud(AR_b)) # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]] +reveal_type(np.flipud(AR_LIKE_b)) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.eye(10)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] +reveal_type(np.eye(10, M=20, dtype=np.int64)) # E: numpy.ndarray[Any, numpy.dtype[{int64}]] +reveal_type(np.eye(10, k=2, dtype=int)) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.diag(AR_b)) # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]] +reveal_type(np.diag(AR_LIKE_b, k=0)) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.diagflat(AR_b)) # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]] +reveal_type(np.diagflat(AR_LIKE_b, k=0)) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.tri(10)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] +reveal_type(np.tri(10, M=20, dtype=np.int64)) # E: numpy.ndarray[Any, numpy.dtype[{int64}]] +reveal_type(np.tri(10, k=2, dtype=int)) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.tril(AR_b)) # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]] +reveal_type(np.tril(AR_LIKE_b, k=0)) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.triu(AR_b)) # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]] +reveal_type(np.triu(AR_LIKE_b, k=0)) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.vander(AR_b)) # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]] +reveal_type(np.vander(AR_u)) # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]] +reveal_type(np.vander(AR_i, N=2)) # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]] +reveal_type(np.vander(AR_f, increasing=True)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] +reveal_type(np.vander(AR_c)) # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]] +reveal_type(np.vander(AR_O)) # E: numpy.ndarray[Any, numpy.dtype[numpy.object_]] + +reveal_type(np.histogram2d(AR_i, AR_b)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]] +reveal_type(np.histogram2d(AR_f, AR_f)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]] +reveal_type(np.histogram2d(AR_f, AR_c, weights=AR_LIKE_b)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]] + +reveal_type(np.mask_indices(10, func1)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{intp}]], numpy.ndarray[Any, numpy.dtype[{intp}]]] +reveal_type(np.mask_indices(8, func2, "0")) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{intp}]], numpy.ndarray[Any, numpy.dtype[{intp}]]] + +reveal_type(np.tril_indices(10)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{int_}]], numpy.ndarray[Any, numpy.dtype[{int_}]]] + +reveal_type(np.tril_indices_from(AR_b)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{int_}]], numpy.ndarray[Any, numpy.dtype[{int_}]]] + +reveal_type(np.triu_indices(10)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{int_}]], numpy.ndarray[Any, numpy.dtype[{int_}]]] + +reveal_type(np.triu_indices_from(AR_b)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{int_}]], numpy.ndarray[Any, numpy.dtype[{int_}]]] |