Skip to content

Commit f782790

Browse files
jorenhamcharris
authored andcommitted
TYP: Fix overlapping overloads issue in 2->1 ufuncs
1 parent 32b58cd commit f782790

File tree

1 file changed

+111
-49
lines changed

1 file changed

+111
-49
lines changed

numpy/_typing/_ufunc.pyi

+111-49
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]
@@ -61,6 +61,13 @@ class _SupportsArrayUFunc(Protocol):
6161
**kwargs: Any,
6262
) -> Any: ...
6363

64+
@type_check_only
65+
class _UFunc3Kwargs(TypedDict, total=False):
66+
where: _ArrayLikeBool_co | None
67+
casting: _CastingKind
68+
order: _OrderKACF
69+
subok: bool
70+
signature: _3Tuple[str | None] | str | None
6471

6572
# NOTE: `reduce`, `accumulate`, `reduceat` and `outer` raise a ValueError for
6673
# ufuncs that don't accept two input arguments and return one output argument.
@@ -72,6 +79,8 @@ class _SupportsArrayUFunc(Protocol):
7279
# NOTE: If 2 output types are returned then `out` must be a
7380
# 2-tuple of arrays. Otherwise `None` or a plain array are also acceptable
7481

82+
# pyright: reportIncompatibleMethodOverride=false
83+
7584
@type_check_only
7685
class _UFunc_Nin1_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc]
7786
@property
@@ -162,34 +171,61 @@ class _UFunc_Nin2_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: i
162171
@property
163172
def signature(self) -> None: ...
164173

165-
@overload
174+
@overload # (scalar, scalar) -> scalar
166175
def __call__(
167176
self,
168-
__x1: _ScalarLike_co,
169-
__x2: _ScalarLike_co,
170-
out: None = ...,
177+
x1: _ScalarLike_co,
178+
x2: _ScalarLike_co,
179+
/,
180+
out: None = None,
171181
*,
172-
where: None | _ArrayLikeBool_co = ...,
173-
casting: _CastingKind = ...,
174-
order: _OrderKACF = ...,
175-
dtype: DTypeLike = ...,
176-
subok: bool = ...,
177-
signature: str | _3Tuple[None | str] = ...,
182+
dtype: DTypeLike | None = None,
183+
**kwds: Unpack[_UFunc3Kwargs],
178184
) -> Any: ...
179-
@overload
185+
@overload # (array-like, array) -> array
180186
def __call__(
181187
self,
182-
__x1: ArrayLike,
183-
__x2: ArrayLike,
184-
out: None | NDArray[Any] | tuple[NDArray[Any]] = ...,
188+
x1: ArrayLike,
189+
x2: NDArray[np.generic],
190+
/,
191+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
185192
*,
186-
where: None | _ArrayLikeBool_co = ...,
187-
casting: _CastingKind = ...,
188-
order: _OrderKACF = ...,
189-
dtype: DTypeLike = ...,
190-
subok: bool = ...,
191-
signature: str | _3Tuple[None | str] = ...,
193+
dtype: DTypeLike | None = None,
194+
**kwds: Unpack[_UFunc3Kwargs],
195+
) -> NDArray[Any]: ...
196+
@overload # (array, array-like) -> array
197+
def __call__(
198+
self,
199+
x1: NDArray[np.generic],
200+
x2: ArrayLike,
201+
/,
202+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
203+
*,
204+
dtype: DTypeLike | None = None,
205+
**kwds: Unpack[_UFunc3Kwargs],
206+
) -> NDArray[Any]: ...
207+
@overload # (array-like, array-like, out=array) -> array
208+
def __call__(
209+
self,
210+
x1: ArrayLike,
211+
x2: ArrayLike,
212+
/,
213+
out: NDArray[np.generic] | tuple[NDArray[np.generic]],
214+
*,
215+
dtype: DTypeLike | None = None,
216+
**kwds: Unpack[_UFunc3Kwargs],
192217
) -> NDArray[Any]: ...
218+
@overload # (array-like, array-like) -> array | scalar
219+
def __call__(
220+
self,
221+
x1: ArrayLike,
222+
x2: ArrayLike,
223+
/,
224+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
225+
*,
226+
dtype: DTypeLike | None = None,
227+
**kwds: Unpack[_UFunc3Kwargs],
228+
) -> NDArray[Any] | Any: ...
193229

194230
def at(
195231
self,
@@ -227,35 +263,61 @@ class _UFunc_Nin2_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: i
227263
out: None | NDArray[Any] = ...,
228264
) -> NDArray[Any]: ...
229265

230-
# Expand `**kwargs` into explicit keyword-only arguments
231-
@overload
266+
@overload # (scalar, scalar) -> scalar
232267
def outer(
233268
self,
234269
A: _ScalarLike_co,
235270
B: _ScalarLike_co,
236-
/, *,
237-
out: None = ...,
238-
where: None | _ArrayLikeBool_co = ...,
239-
casting: _CastingKind = ...,
240-
order: _OrderKACF = ...,
241-
dtype: DTypeLike = ...,
242-
subok: bool = ...,
243-
signature: str | _3Tuple[None | str] = ...,
271+
/,
272+
*,
273+
out: None = None,
274+
dtype: DTypeLike | None = None,
275+
**kwds: Unpack[_UFunc3Kwargs],
244276
) -> Any: ...
245-
@overload
246-
def outer( # type: ignore[misc]
277+
@overload # (array-like, array) -> array
278+
def outer(
247279
self,
248280
A: ArrayLike,
281+
B: NDArray[np.generic],
282+
/,
283+
*,
284+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
285+
dtype: DTypeLike | None = None,
286+
**kwds: Unpack[_UFunc3Kwargs],
287+
) -> NDArray[Any]: ...
288+
@overload # (array, array-like) -> array
289+
def outer(
290+
self,
291+
A: NDArray[np.generic],
249292
B: ArrayLike,
250-
/, *,
251-
out: None | NDArray[Any] | tuple[NDArray[Any]] = ...,
252-
where: None | _ArrayLikeBool_co = ...,
253-
casting: _CastingKind = ...,
254-
order: _OrderKACF = ...,
255-
dtype: DTypeLike = ...,
256-
subok: bool = ...,
257-
signature: str | _3Tuple[None | str] = ...,
293+
/,
294+
*,
295+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
296+
dtype: DTypeLike | None = None,
297+
**kwds: Unpack[_UFunc3Kwargs],
258298
) -> NDArray[Any]: ...
299+
@overload # (array-like, array-like, out=array) -> array
300+
def outer(
301+
self,
302+
A: ArrayLike,
303+
B: ArrayLike,
304+
/,
305+
*,
306+
out: NDArray[np.generic] | tuple[NDArray[np.generic]],
307+
dtype: DTypeLike | None = None,
308+
**kwds: Unpack[_UFunc3Kwargs],
309+
) -> NDArray[Any]: ...
310+
@overload # (array-like, array-like) -> array | scalar
311+
def outer(
312+
self,
313+
A: ArrayLike,
314+
B: ArrayLike,
315+
/,
316+
*,
317+
out: NDArray[np.generic] | tuple[NDArray[np.generic]] | None = None,
318+
dtype: DTypeLike | None = None,
319+
**kwds: Unpack[_UFunc3Kwargs],
320+
) -> NDArray[Any] | Any: ...
259321

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

0 commit comments

Comments
 (0)