diff options
author | Bas van Beek <43369155+BvB93@users.noreply.github.com> | 2022-01-11 18:52:42 +0100 |
---|---|---|
committer | Bas van Beek <43369155+BvB93@users.noreply.github.com> | 2022-01-11 18:56:27 +0100 |
commit | fc30b217c1021b268a8cdcc594fa9d792ce17b19 (patch) | |
tree | dffd18d3121f58070e72c823e114883e9fd04832 /numpy | |
parent | d2df8c6a831a9a4c297c4e423da1292ecd8228f4 (diff) | |
download | numpy-fc30b217c1021b268a8cdcc594fa9d792ce17b19.tar.gz |
ENH: Type the `like` parameter via a `__array_function__` protocol
Diffstat (limited to 'numpy')
-rw-r--r-- | numpy/core/_asarray.pyi | 8 | ||||
-rw-r--r-- | numpy/core/multiarray.pyi | 101 | ||||
-rw-r--r-- | numpy/core/numeric.pyi | 21 | ||||
-rw-r--r-- | numpy/lib/npyio.pyi | 20 | ||||
-rw-r--r-- | numpy/lib/twodim_base.pyi | 13 | ||||
-rw-r--r-- | numpy/typing/tests/data/fail/array_constructors.pyi | 2 | ||||
-rw-r--r-- | numpy/typing/tests/data/reveal/array_constructors.pyi | 1 |
7 files changed, 89 insertions, 77 deletions
diff --git a/numpy/core/_asarray.pyi b/numpy/core/_asarray.pyi index 0da2de912..51b794130 100644 --- a/numpy/core/_asarray.pyi +++ b/numpy/core/_asarray.pyi @@ -2,7 +2,7 @@ from collections.abc import Iterable from typing import TypeVar, Union, overload, Literal from numpy import ndarray -from numpy.typing import ArrayLike, DTypeLike +from numpy.typing import DTypeLike, _SupportsArrayFunc _ArrayType = TypeVar("_ArrayType", bound=ndarray) @@ -22,7 +22,7 @@ def require( dtype: None = ..., requirements: None | _Requirements | Iterable[_Requirements] = ..., *, - like: ArrayLike = ... + like: _SupportsArrayFunc = ... ) -> _ArrayType: ... @overload def require( @@ -30,7 +30,7 @@ def require( dtype: DTypeLike = ..., requirements: _E | Iterable[_RequirementsWithE] = ..., *, - like: ArrayLike = ... + like: _SupportsArrayFunc = ... ) -> ndarray: ... @overload def require( @@ -38,5 +38,5 @@ def require( dtype: DTypeLike = ..., requirements: None | _Requirements | Iterable[_Requirements] = ..., *, - like: ArrayLike = ... + like: _SupportsArrayFunc = ... ) -> ndarray: ... diff --git a/numpy/core/multiarray.pyi b/numpy/core/multiarray.pyi index 423aed85e..f2d3622d2 100644 --- a/numpy/core/multiarray.pyi +++ b/numpy/core/multiarray.pyi @@ -61,6 +61,7 @@ from numpy.typing import ( NDArray, ArrayLike, _SupportsArray, + _SupportsArrayFunc, _NestedSequence, _FiniteNestedSequence, _ArrayLikeBool_co, @@ -177,7 +178,7 @@ def array( order: _OrderKACF = ..., subok: L[True], ndmin: int = ..., - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> _ArrayType: ... @overload def array( @@ -188,7 +189,7 @@ def array( order: _OrderKACF = ..., subok: bool = ..., ndmin: int = ..., - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def array( @@ -199,7 +200,7 @@ def array( order: _OrderKACF = ..., subok: bool = ..., ndmin: int = ..., - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload def array( @@ -210,7 +211,7 @@ def array( order: _OrderKACF = ..., subok: bool = ..., ndmin: int = ..., - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def array( @@ -221,7 +222,7 @@ def array( order: _OrderKACF = ..., subok: bool = ..., ndmin: int = ..., - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -230,7 +231,7 @@ def zeros( dtype: None = ..., order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload def zeros( @@ -238,7 +239,7 @@ def zeros( dtype: _DTypeLike[_SCT], order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def zeros( @@ -246,7 +247,7 @@ def zeros( dtype: DTypeLike, order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -255,7 +256,7 @@ def empty( dtype: None = ..., order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload def empty( @@ -263,7 +264,7 @@ def empty( dtype: _DTypeLike[_SCT], order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def empty( @@ -271,7 +272,7 @@ def empty( dtype: DTypeLike, order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -468,7 +469,7 @@ def asarray( dtype: None = ..., order: _OrderKACF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def asarray( @@ -476,7 +477,7 @@ def asarray( dtype: None = ..., order: _OrderKACF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload def asarray( @@ -484,7 +485,7 @@ def asarray( dtype: _DTypeLike[_SCT], order: _OrderKACF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def asarray( @@ -492,7 +493,7 @@ def asarray( dtype: DTypeLike, order: _OrderKACF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -501,7 +502,7 @@ def asanyarray( dtype: None = ..., order: _OrderKACF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> _ArrayType: ... @overload def asanyarray( @@ -509,7 +510,7 @@ def asanyarray( dtype: None = ..., order: _OrderKACF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def asanyarray( @@ -517,7 +518,7 @@ def asanyarray( dtype: None = ..., order: _OrderKACF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload def asanyarray( @@ -525,7 +526,7 @@ def asanyarray( dtype: _DTypeLike[_SCT], order: _OrderKACF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def asanyarray( @@ -533,7 +534,7 @@ def asanyarray( dtype: DTypeLike, order: _OrderKACF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -541,28 +542,28 @@ def ascontiguousarray( a: _ArrayLike[_SCT], dtype: None = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def ascontiguousarray( a: object, dtype: None = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload def ascontiguousarray( a: Any, dtype: _DTypeLike[_SCT], *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def ascontiguousarray( a: Any, dtype: DTypeLike, *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -570,28 +571,28 @@ def asfortranarray( a: _ArrayLike[_SCT], dtype: None = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def asfortranarray( a: object, dtype: None = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload def asfortranarray( a: Any, dtype: _DTypeLike[_SCT], *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def asfortranarray( a: Any, dtype: DTypeLike, *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... # In practice `list[Any]` is list with an int, int and a valid @@ -609,7 +610,7 @@ def fromstring( count: SupportsIndex = ..., *, sep: str, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload def fromstring( @@ -618,7 +619,7 @@ def fromstring( count: SupportsIndex = ..., *, sep: str, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def fromstring( @@ -627,7 +628,7 @@ def fromstring( count: SupportsIndex = ..., *, sep: str, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... def frompyfunc( @@ -646,7 +647,7 @@ def fromfile( sep: str = ..., offset: SupportsIndex = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload def fromfile( @@ -656,7 +657,7 @@ def fromfile( sep: str = ..., offset: SupportsIndex = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def fromfile( @@ -666,7 +667,7 @@ def fromfile( sep: str = ..., offset: SupportsIndex = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -675,7 +676,7 @@ def fromiter( dtype: _DTypeLike[_SCT], count: SupportsIndex = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def fromiter( @@ -683,7 +684,7 @@ def fromiter( dtype: DTypeLike, count: SupportsIndex = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -693,7 +694,7 @@ def frombuffer( count: SupportsIndex = ..., offset: SupportsIndex = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload def frombuffer( @@ -702,7 +703,7 @@ def frombuffer( count: SupportsIndex = ..., offset: SupportsIndex = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def frombuffer( @@ -711,7 +712,7 @@ def frombuffer( count: SupportsIndex = ..., offset: SupportsIndex = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -719,7 +720,7 @@ def arange( # type: ignore[misc] stop: _IntLike_co, /, *, dtype: None = ..., - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[signedinteger[Any]]: ... @overload def arange( # type: ignore[misc] @@ -728,14 +729,14 @@ def arange( # type: ignore[misc] step: _IntLike_co = ..., dtype: None = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[signedinteger[Any]]: ... @overload def arange( # type: ignore[misc] stop: _FloatLike_co, /, *, dtype: None = ..., - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[floating[Any]]: ... @overload def arange( # type: ignore[misc] @@ -744,14 +745,14 @@ def arange( # type: ignore[misc] step: _FloatLike_co = ..., dtype: None = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[floating[Any]]: ... @overload def arange( stop: _TD64Like_co, /, *, dtype: None = ..., - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[timedelta64]: ... @overload def arange( @@ -760,7 +761,7 @@ def arange( step: _TD64Like_co = ..., dtype: None = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[timedelta64]: ... @overload def arange( # both start and stop must always be specified for datetime64 @@ -769,14 +770,14 @@ def arange( # both start and stop must always be specified for datetime64 step: datetime64 = ..., dtype: None = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[datetime64]: ... @overload def arange( stop: Any, /, *, dtype: _DTypeLike[_SCT], - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def arange( @@ -785,14 +786,14 @@ def arange( step: Any = ..., dtype: _DTypeLike[_SCT] = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def arange( stop: Any, /, *, dtype: DTypeLike, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload def arange( @@ -801,7 +802,7 @@ def arange( step: Any = ..., dtype: DTypeLike = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... def datetime_data( diff --git a/numpy/core/numeric.pyi b/numpy/core/numeric.pyi index 8b92abab4..d5e28d24c 100644 --- a/numpy/core/numeric.pyi +++ b/numpy/core/numeric.pyi @@ -37,6 +37,7 @@ from numpy.typing import ( _SupportsDType, _FiniteNestedSequence, _SupportsArray, + _SupportsArrayFunc, _ScalarLike_co, _ArrayLikeBool_co, _ArrayLikeUInt_co, @@ -108,7 +109,7 @@ def ones( dtype: None = ..., order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload def ones( @@ -116,7 +117,7 @@ def ones( dtype: _DTypeLike[_SCT], order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def ones( @@ -124,7 +125,7 @@ def ones( dtype: DTypeLike, order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -175,7 +176,7 @@ def full( dtype: None = ..., order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload def full( @@ -184,7 +185,7 @@ def full( dtype: _DTypeLike[_SCT], order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def full( @@ -193,7 +194,7 @@ def full( dtype: DTypeLike, order: _OrderCF = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -563,7 +564,7 @@ def fromfunction( shape: Sequence[int], *, dtype: DTypeLike = ..., - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., **kwargs: Any, ) -> _T: ... @@ -584,21 +585,21 @@ def identity( n: int, dtype: None = ..., *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload def identity( n: int, dtype: _DTypeLike[_SCT], *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def identity( n: int, dtype: DTypeLike, *, - like: ArrayLike = ..., + like: _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... def allclose( diff --git a/numpy/lib/npyio.pyi b/numpy/lib/npyio.pyi index 60684c846..4f11fa807 100644 --- a/numpy/lib/npyio.pyi +++ b/numpy/lib/npyio.pyi @@ -27,7 +27,13 @@ from numpy import ( ) from numpy.ma.mrecords import MaskedRecords -from numpy.typing import ArrayLike, DTypeLike, NDArray, _SupportsDType +from numpy.typing import ( + ArrayLike, + DTypeLike, + NDArray, + _SupportsDType, + _SupportsArrayFunc, +) from numpy.core.multiarray import ( packbits as packbits, @@ -144,7 +150,7 @@ def loadtxt( encoding: None | str = ..., max_rows: None | int = ..., *, - like: None | ArrayLike = ... + like: None | _SupportsArrayFunc = ... ) -> NDArray[float64]: ... @overload def loadtxt( @@ -160,7 +166,7 @@ def loadtxt( encoding: None | str = ..., max_rows: None | int = ..., *, - like: None | ArrayLike = ... + like: None | _SupportsArrayFunc = ... ) -> NDArray[_SCT]: ... @overload def loadtxt( @@ -176,7 +182,7 @@ def loadtxt( encoding: None | str = ..., max_rows: None | int = ..., *, - like: None | ArrayLike = ... + like: None | _SupportsArrayFunc = ... ) -> NDArray[Any]: ... def savetxt( @@ -233,7 +239,7 @@ def genfromtxt( encoding: str = ..., *, ndmin: L[0, 1, 2] = ..., - like: None | ArrayLike = ..., + like: None | _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload def genfromtxt( @@ -262,7 +268,7 @@ def genfromtxt( encoding: str = ..., *, ndmin: L[0, 1, 2] = ..., - like: None | ArrayLike = ..., + like: None | _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def genfromtxt( @@ -291,7 +297,7 @@ def genfromtxt( encoding: str = ..., *, ndmin: L[0, 1, 2] = ..., - like: None | ArrayLike = ..., + like: None | _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload diff --git a/numpy/lib/twodim_base.pyi b/numpy/lib/twodim_base.pyi index 76d7e5a9d..491862408 100644 --- a/numpy/lib/twodim_base.pyi +++ b/numpy/lib/twodim_base.pyi @@ -31,6 +31,7 @@ from numpy.typing import ( NDArray, _FiniteNestedSequence, _SupportsArray, + _SupportsArrayFunc, _ArrayLikeInt_co, _ArrayLikeFloat_co, _ArrayLikeComplex_co, @@ -73,7 +74,7 @@ def eye( dtype: None = ..., order: _OrderCF = ..., *, - like: None | ArrayLike = ..., + like: None | _SupportsArrayFunc = ..., ) -> NDArray[float64]: ... @overload def eye( @@ -83,7 +84,7 @@ def eye( dtype: _DTypeLike[_SCT] = ..., order: _OrderCF = ..., *, - like: None | ArrayLike = ..., + like: None | _SupportsArrayFunc = ..., ) -> NDArray[_SCT]: ... @overload def eye( @@ -93,7 +94,7 @@ def eye( dtype: DTypeLike = ..., order: _OrderCF = ..., *, - like: None | ArrayLike = ..., + like: None | _SupportsArrayFunc = ..., ) -> NDArray[Any]: ... @overload @@ -113,7 +114,7 @@ def tri( k: int = ..., dtype: None = ..., *, - like: None | ArrayLike = ... + like: None | _SupportsArrayFunc = ... ) -> NDArray[float64]: ... @overload def tri( @@ -122,7 +123,7 @@ def tri( k: int = ..., dtype: _DTypeLike[_SCT] = ..., *, - like: None | ArrayLike = ... + like: None | _SupportsArrayFunc = ... ) -> NDArray[_SCT]: ... @overload def tri( @@ -131,7 +132,7 @@ def tri( k: int = ..., dtype: DTypeLike = ..., *, - like: None | ArrayLike = ... + like: None | _SupportsArrayFunc = ... ) -> NDArray[Any]: ... @overload diff --git a/numpy/typing/tests/data/fail/array_constructors.pyi b/numpy/typing/tests/data/fail/array_constructors.pyi index 065b7d8a0..278894631 100644 --- a/numpy/typing/tests/data/fail/array_constructors.pyi +++ b/numpy/typing/tests/data/fail/array_constructors.pyi @@ -29,3 +29,5 @@ np.geomspace(None, 'bob') # E: No overload variant np.stack(generator) # E: No overload variant np.hstack({1, 2}) # E: No overload variant np.vstack(1) # E: No overload variant + +np.array([1], like=1) # 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 c3b0c3457..448ed7e8b 100644 --- a/numpy/typing/tests/data/reveal/array_constructors.pyi +++ b/numpy/typing/tests/data/reveal/array_constructors.pyi @@ -28,6 +28,7 @@ reveal_type(np.array(B, subok=True)) # E: SubClass[{float64}] reveal_type(np.array([1, 1.0])) # E: ndarray[Any, dtype[Any]] reveal_type(np.array(A, dtype=np.int64)) # E: ndarray[Any, dtype[{int64}]] reveal_type(np.array(A, dtype='c16')) # E: ndarray[Any, dtype[Any]] +reveal_type(np.array(A, like=A)) # E: ndarray[Any, dtype[{float64}]] reveal_type(np.zeros([1, 5, 6])) # E: ndarray[Any, dtype[{float64}]] reveal_type(np.zeros([1, 5, 6], dtype=np.int64)) # E: ndarray[Any, dtype[{int64}]] |