8000 TYP: Fix overlapping overloads issue in 2->1 ufuncs · melissawm/numpy@86eea0c · GitHub
[go: up one dir, main page]
Skip to content

Commit 86eea0c

Browse files
jorenhamjorenham
committed
TYP: Fix overlapping overloads issue in 2->1 ufuncs
1 parent 1e10174 commit 86eea0c

File tree

1 file changed

+112
-49
lines changed

1 file changed

+112
-49
lines changed

‎numpy/_typing/_ufunc.pyi

Lines changed: 112 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -4,32 +4,32 @@ The signatures of the ufuncs are too varied to reasonably type
44
with a single class. So instead, `ufunc` has been expanded into
55
four private subclasses, one for each combination of
66
`~ufunc.nin` and `~ufunc.nout`.
7-
87
"""
98

109
from typing import (
1110
Any,
1211
Generic,
12+
Literal,
1313
NoReturn,
14-
TypedDict,
15-
overload,
14+
Protocol,
15+
SupportsIndex,
1616
TypeAlias,
17+
TypedDict,
1718
TypeVar,
18-
Literal,
19-
SupportsIndex,
20-
Protocol,
19+
overload,
2120
type_check_only,
2221
)
22+
2323
from typing_extensions import LiteralString, Unpack
2424

2525
import numpy as np
26-
from numpy import ufunc, _CastingKind, _OrderKACF
26+
from numpy import _CastingKind, _OrderKACF, ufunc
2727
from numpy.typing import NDArray
2828

29-
from ._shape import _ShapeLike
30-
from ._scalars import _ScalarLike_co
3129
from ._array_like import ArrayLike, _ArrayLikeBool_co, _ArrayLikeInt_co
3230
from ._dtype_like import DTypeLike
31+
from ._scalars import _ScalarLike_co
32+
from ._shape import _ShapeLike
3333

3434
_T = TypeVar("_T")
3535
_2Tuple: TypeAlias = tuple[_T, _T]
@@ -60,6 +60,14 @@ class _SupportsArrayUFunc(Protocol):
6060
**kwargs: Any,
6161
) -> Any: ...
6262

63+
@type_check_only
64+
class _UFunc3Kwargs(TypedDict, total=False):
65+
where: _ArrayLikeBool_co | None
66+
casting: _CastingKind
67+
order: _OrderKACF
68+
subok: bool
69+
signature: _3Tuple[str | None] | str | None
70+
6371
# NOTE: `reduce`, `accumulate`, `reduceat` and `outer` raise a ValueError for
6472
# ufuncs that don't accept two input arguments and return one output argument.
6573
# In such cases the respective methods return `NoReturn`
@@ -70,6 +78,8 @@ class _SupportsArrayUFunc(Protocol):
7078
# NOTE: If 2 output types are returned then `out` must be a
7179
# 2-tuple of arrays. Otherwise `None` or a plain array are also acceptable
7280

81+
# pyright: reportIncompatibleMethodOverride=false
82+
7383
@type_check_only
7484
class _UFunc_Nin1_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc]
7585
@property
@@ -160,34 +170,61 @@ class _UFunc_Nin2_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: i
160170
@property
161171
def signature(self) -> None: ...
162172

163-
@overload
173+
@overload # (scalar, scalar) -> scalar
164174
def __call__(
165175
self,
166-
__x1: _ScalarLike_co,
167-
__x2: _ScalarLike_co,
168-
out: None = ...,
176+
x1: _ScalarLike_co,
177+
x2: _ScalarLike_co,
178+
/,
179+
out: None = None,
169180
*,
170-
where: None | _ArrayLikeBool_co = ...,
171-
casting: _CastingKind = ...,
172-
order: _OrderKACF = ...,
173-
dtype: DTypeLike = ...,
174-
subok: bool = ...,
175-
signature: str | _3Tuple[None | str] = ...,
181+
dtype: DTypeLike | None = None,
182+
**kwds: Unpack[_UFunc3Kwargs],
176183
) -> Any: ...
177-
@overload
184+
@overload # (array-like, array) -> array
178185
def __call__(
179186
self,
180-
__x1: ArrayLike,
181-
__x2: ArrayLike,
182-
out: None | NDArray[Any] | tuple[NDArray[Any]] = ...,
187+
x1: ArrayLike,
188+
x2: NDArray[np.generic],
189+
/,
190+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
183191
*,
184-
where: None | _ArrayLikeBool_co = ...,
185-
casting: _CastingKind = ...,
186-
order: _OrderKACF = ...,
187-
dtype: DTypeLike = ...,
188-
subok: bool = ...,
189-
signature: str | _3Tuple[None | str] = ...,
192+
dtype: DTypeLike | None = None,
193+
**kwds: Unpack[_UFunc3Kwargs],
190194
) -> NDArray[Any]: ...
195+
@overload # (array, array-like) -> array
196+
def __call__(
197+
self,
198+
x1: NDArray[np.generic],
199+
x2: ArrayLike,
200+
/,
201+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
202+
*,
203+
dtype: DTypeLike | None = None,
204+
**kwds: Unpack[_UFunc3Kwargs],
205+
) -> NDArray[Any]: ...
206+
@overload # (array-like, array-like, out=array) -> array
207+
def __call__(
208+
self,
209+
x1: ArrayLike,
210+
x2: ArrayLike,
211+
/,
212+
out: NDArray[np.generic] | tuple[NDArray[np.generic]],
213+
*,
214+
dtype: DTypeLike | None = None,
215+
**kwds: Unpack[_UFunc3Kwargs],
216+
) -> NDArray[Any]: ...
217+
@overload # (array-like, array-like) -> array | scalar
218+
def __call__(
219+
self,
220+
x1: ArrayLike,
221+
x2: ArrayLike,
222+
/,
223+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
224+
*,
225+
dtype: DTypeLike | None = None,
226+
**kwds: Unpack[_UFunc3Kwargs],
227+
) -> NDArray[Any] | Any: ...
191228

192229
def at(
193230
self,
@@ -225,35 +262,61 @@ class _UFunc_Nin2_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: i
225262
out: None | NDArray[Any] = ...,
226263
) -> NDArray[Any]: ...
227264

228-
# Expand `**kwargs` into explicit keyword-only arguments
229-
@overload
265+
@overload # (scalar, scalar) -> scalar
230266
def outer(
231267
self,
232268
A: _ScalarLike_co,
233269
B: _ScalarLike_co,
234-
/, *,
235-
out: None = ...,
236-
where: None | _ArrayLikeBool_co = ...,
237-
casting: _CastingKind = ...,
238-
order: _OrderKACF = ...,
239-
dtype: DTypeLike = ...,
240-
subok: bool = ...,
241-
signature: str | _3Tuple[None | str] = ...,
270+
/,
271+
*,
272+
out: None = None,
273+
dtype: DTypeLike | None = None,
274+
**kwds: Unpack[_UFunc3Kwargs],
242275
) -> Any: ...
243-
@overload
244-
def outer( # type: ignore[misc]
276+
@overload # (array-like, array) -> array
277+
def outer(
245278
self,
246279
A: ArrayLike,
280+
B: NDArray[np.generic],
281+
/,
282+
*,
283+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
284+
dtype: DTypeLike | None = None,
285+
**kwds: Unpack[_UFunc3Kwargs],
286+
) -> NDArray[Any]: ...
287+
@overload # (array, array-like) -> array
288+
def outer(
289+
self,
290+
A: NDArray[np.generic],
247291
B: ArrayLike,
248-
/, *,
249-
out: None | NDArray[Any] | tuple[NDArray[Any]] = ...,
250-
where: None | _ArrayLikeBool_co = ...,
251-
casting: _CastingKind = ...,
252-
order: _OrderKACF = ...,
253-
dtype: DTypeLike = ...,
254-
subok: bool = ...,
255-
signature: str | _3Tuple[None | str] = ...,
292+
/,
293+
*,
294+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
295+
dtype: DTypeLike | None = None,
296+
**kwds: Unpack[_UFunc3Kwargs],
256297
) -> NDArray[Any]: ...
298+
@overload # (array-like, array-like, out=array) -> array
299+
def outer(
300+
self,
301+
A: ArrayLike,
302+
B: ArrayLike,
303+
/,
304+
*,
305+
out: NDArray[np.generic] | tuple[NDArray[np.generic]],
306+
dtype: DTypeLike | None = None,
307+
**kwds: Unpack[_UFunc3Kwargs],
308+
) -> NDArray[Any]: ...
309+
@overload # (array-like, array-like) -> array | scalar
310+
def outer(
311+
self,
312+
A: ArrayLike,
313+
B: ArrayLike,
314+
/,
6312 315+
*,
316+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
317+
dtype: DTypeLike | None = None,
318+
**kwds: Unpack[_UFunc3Kwargs],
319+
) -> NDArray[Any] | Any: ...
257320

258321
@type_check_only
259322
class _UFunc_Nin1_Nout2(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc]

0 commit comments

Comments
 (0)
0