@@ -4,32 +4,32 @@ The signatures of the ufuncs are too varied to reasonably type
4
4
with a single class. So instead, `ufunc` has been expanded into
5
5
four private subclasses, one for each combination of
6
6
`~ufunc.nin` and `~ufunc.nout`.
7
-
8
7
"""
9
8
10
9
from typing import (
11
10
Any ,
12
11
Generic ,
12
+ Literal ,
13
13
NoReturn ,
14
- TypedDict ,
15
- overload ,
14
+ Protocol ,
15
+ SupportsIndex ,
16
16
TypeAlias ,
17
+ TypedDict ,
17
18
TypeVar ,
18
- Literal ,
19
- SupportsIndex ,
20
- Protocol ,
19
+ overload ,
21
20
type_check_only ,
22
21
)
22
+
23
23
from typing_extensions import LiteralString , Unpack
24
24
25
25
import numpy as np
26
- from numpy import ufunc , _CastingKind , _OrderKACF
26
+ from numpy import _CastingKind , _OrderKACF , ufunc
27
27
from numpy .typing import NDArray
28
28
29
- from ._shape import _ShapeLike
30
- from ._scalars import _ScalarLike_co
31
29
from ._array_like import ArrayLike , _ArrayLikeBool_co , _ArrayLikeInt_co
32
30
from ._dtype_like import DTypeLike
31
+ from ._scalars import _ScalarLike_co
32
+ from ._shape import _ShapeLike
33
33
34
34
_T = TypeVar ("_T" )
35
35
_2Tuple : TypeAlias = tuple [_T , _T ]
@@ -61,6 +61,13 @@ class _SupportsArrayUFunc(Protocol):
61
61
** kwargs : Any ,
62
62
) -> Any : ...
63
63
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
64
71
65
72
# NOTE: `reduce`, `accumulate`, `reduceat` and `outer` raise a ValueError for
66
73
# ufuncs that don't accept two input arguments and return one output argument.
@@ -72,6 +79,8 @@ class _SupportsArrayUFunc(Protocol):
72
79
# NOTE: If 2 output types are returned then `out` must be a
73
80
# 2-tuple of arrays. Otherwise `None` or a plain array are also acceptable
74
81
82
+ # pyright: reportIncompatibleMethodOverride=false
83
+
75
84
@type_check_only
76
85
class _UFunc_Nin1_Nout1 (ufunc , Generic [_NameType , _NTypes , _IDType ]): # type: ignore[misc]
77
86
@property
@@ -162,34 +171,61 @@ class _UFunc_Nin2_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: i
162
171
@property
163
172
def signature (self ) -> None : ...
164
173
165
- @overload
174
+ @overload # (scalar, scalar) -> scalar
166
175
def __call__ (
167
176
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 ,
171
181
* ,
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 ],
178
184
) -> Any : ...
179
- @overload
185
+ @overload # (array-like, array) -> array
180
186
def __call__ (
181
187
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 ,
185
192
* ,
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 ],
192
217
) -> 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 : ...
193
229
194
230
def at (
195
231
self ,
@@ -227,35 +263,61 @@ class _UFunc_Nin2_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: i
227
263
out : None | NDArray [Any ] = ...,
228
264
) -> NDArray [Any ]: ...
229
265
230
- # Expand `**kwargs` into explicit keyword-only arguments
231
- @overload
266
+ @overload # (scalar, scalar) -> scalar
232
267
def outer (
233
268
self ,
234
269
A : _ScalarLike_co ,
235
270
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 ],
244
276
) -> Any : ...
245
- @overload
246
- def outer ( # type: ignore[misc]
277
+ @overload # (array-like, array) -> array
278
+ def outer (
247
279
self ,
248
280
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 ],
249
292
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 ],
258
298
) -> 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 : ...
259
321
260
322
@type_check_only
261
323
class _UFunc_Nin1_Nout2 (ufunc , Generic [_NameType , _NTypes , _IDType ]): # type: ignore[misc]
0 commit comments