Skip to content

Commit 3358464

Browse files
authored
Merge pull request #680 from numpy/ndarray-binop-typevar-fix
2 parents daae953 + 46e704f commit 3358464

1 file changed

Lines changed: 36 additions & 19 deletions

File tree

src/numpy-stubs/__init__.pyi

Lines changed: 36 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -611,6 +611,8 @@ _AnyShapeT = TypeVar(
611611
_nt.Shape1N,
612612
_nt.Shape0N,
613613
)
614+
_AnyItemT = TypeVar("_AnyItemT", bool, int, float, complex, bytes, str, dt.datetime, dt.date, dt.timedelta)
615+
_AnyNumberItemT = TypeVar("_AnyNumberItemT", int, float, complex)
614616

615617
###
616618
# Type parameters (for internal use only)
@@ -654,7 +656,6 @@ _CharDTypeT = TypeVar("_CharDTypeT", bound=dtype[character])
654656

655657
_ItemT_co = TypeVar("_ItemT_co", default=Any, covariant=True)
656658
_BoolItemT_co = TypeVar("_BoolItemT_co", bound=py_bool, default=py_bool, covariant=True)
657-
_NumberItemT = TypeVar("_NumberItemT", bound=complex)
658659
_NumberItemT_co = TypeVar("_NumberItemT_co", bound=complex, default=Any, covariant=True)
659660
_InexactItemT_co = TypeVar("_InexactItemT_co", bound=complex, default=Any, covariant=True)
660661
_FlexItemT_co = TypeVar(
@@ -1932,7 +1933,9 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
19321933
@overload
19331934
def __add__(self: _nt.StringArrayND[_T], x: _nt.ToString_nd[_T] | _nt.ToStr_nd, /) -> _nt.StringArrayND[_T]: ...
19341935
@overload
1935-
def __add__(self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRAdd[_T]], /) -> _nt.Array[Incomplete]: ...
1936+
def __add__(
1937+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanRAdd[_AnyItemT]], /
1938+
) -> _nt.Array[Incomplete]: ...
19361939

19371940
#
19381941
@overload
@@ -1958,7 +1961,9 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
19581961
@overload
19591962
def __radd__(self: _nt.StringArrayND[_T], x: _nt.ToString_nd[_T] | _nt.ToStr_nd, /) -> _nt.StringArrayND[_T]: ...
19601963
@overload
1961-
def __radd__(self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanAdd[_T]], /) -> _nt.Array[Incomplete]: ...
1964+
def __radd__(
1965+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanAdd[_AnyItemT]], /
1966+
) -> _nt.Array[Incomplete]: ...
19621967

19631968
#
19641969
@overload # type: ignore[misc]
@@ -1981,7 +1986,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
19811986
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
19821987
@overload
19831988
def __iadd__(
1984-
self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRAdd[_T, _T]], /
1989+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanRAdd[_AnyItemT, _AnyItemT]], /
19851990
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
19861991

19871992
#
@@ -2005,7 +2010,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
20052010
def __sub__(self: _nt.Array[object_], x: object, /) -> _nt.Array[object_]: ...
20062011
@overload
20072012
def __sub__(
2008-
self: _nt.Array[generic[_NumberItemT]], x: _nt.Sequence1ND[_nt.op.CanRSub[_NumberItemT]], /
2013+
self: _nt.Array[generic[_AnyNumberItemT]], x: _nt.Sequence1ND[_nt.op.CanRSub[_AnyNumberItemT]], /
20092014
) -> _nt.Array[Incomplete]: ...
20102015

20112016
#
@@ -2029,7 +2034,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
20292034
def __rsub__(self: _nt.Array[object_], x: object, /) -> _nt.Array[object_]: ...
20302035
@overload
20312036
def __rsub__(
2032-
self: _nt.Array[generic[_NumberItemT]], x: _nt.Sequence1ND[_nt.op.CanSub[_NumberItemT]], /
2037+
self: _nt.Array[generic[_AnyNumberItemT]], x: _nt.Sequence1ND[_nt.op.CanSub[_AnyNumberItemT]], /
20332038
) -> _nt.Array[Incomplete]: ...
20342039

20352040
#
@@ -2047,7 +2052,9 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
20472052
def __isub__(self: _nt.Array[object_], x: object, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
20482053
@overload
20492054
def __isub__(
2050-
self: _nt.Array[generic[_NumberItemT]], x: _nt.Sequence1ND[_nt.op.CanRSub[_T, _T]], /
2055+
self: _nt.Array[generic[_AnyNumberItemT]],
2056+
x: _nt.Sequence1ND[_nt.op.CanRSub[_AnyNumberItemT, _AnyNumberItemT]],
2057+
/,
20512058
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
20522059

20532060
#
@@ -2072,7 +2079,9 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
20722079
@overload
20732080
def __mul__(self: _nt.StringArrayND[_T], x: _nt.ToInteger_nd, /) -> _nt.StringArrayND[_T]: ...
20742081
@overload
2075-
def __mul__(self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRMul[_T]], /) -> _nt.Array[Incomplete]: ...
2082+
def __mul__(
2083+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanRMul[_AnyItemT]], /
2084+
) -> _nt.Array[Incomplete]: ...
20762085

20772086
#
20782087
@overload
@@ -2096,7 +2105,9 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
20962105
@overload
20972106
def __rmul__(self: _nt.StringArrayND[_T], x: _nt.ToInteger_nd, /) -> _nt.StringArrayND[_T]: ...
20982107
@overload
2099-
def __rmul__(self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanMul[_T]], /) -> _nt.Array[Incomplete]: ...
2108+
def __rmul__(
2109+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanMul[_AnyItemT]], /
2110+
) -> _nt.Array[Incomplete]: ...
21002111

21012112
#
21022113
@overload # type: ignore[misc]
@@ -2117,7 +2128,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
21172128
def __imul__(self: _nt.StringArrayND[_T], x: _nt.ToInteger_nd, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
21182129
@overload
21192130
def __imul__(
2120-
self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRMul[_T, _T]], /
2131+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanRMul[_AnyItemT, _AnyItemT]], /
21212132
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
21222133

21232134
# TODO(jorenham): Support the "1d @ 1d -> scalar" case
@@ -2267,7 +2278,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
22672278
) -> _nt.Array[timedelta64]: ...
22682279
@overload
22692280
def __truediv__(
2270-
self: _nt.Array[generic[_NumberItemT]], x: _nt.Sequence1ND[_nt.op.CanRTruediv[_NumberItemT]], /
2281+
self: _nt.Array[generic[_AnyNumberItemT]], x: _nt.Sequence1ND[_nt.op.CanRTruediv[_AnyNumberItemT]], /
22712282
) -> _nt.Array[Incomplete]: ...
22722283
@overload
22732284
def __truediv__(self: _nt.Array[object_], x: object, /) -> _nt.Array[object_]: ...
@@ -2293,7 +2304,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
22932304
def __rtruediv__(self: _nt.Array[integer | floating], x: _nt.ToTimeDelta_nd, /) -> _nt.Array[timedelta64]: ...
22942305
@overload
22952306
def __rtruediv__(
2296-
self: _nt.Array[generic[_NumberItemT]], x: _nt.Sequence1ND[_nt.op.CanTruediv[_NumberItemT]], /
2307+
self: _nt.Array[generic[_AnyNumberItemT]], x: _nt.Sequence1ND[_nt.op.CanTruediv[_AnyNumberItemT]], /
22972308
) -> _nt.Array[Incomplete]: ...
22982309
@overload
22992310
def __rtruediv__(self: _nt.Array[object_], x: object, /) -> _nt.Array[object_]: ...
@@ -2311,7 +2322,9 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
23112322
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
23122323
@overload
23132324
def __itruediv__(
2314-
self: _nt.Array[generic[_NumberItemT]], x: _nt.Sequence1ND[_nt.op.CanRTruediv[_NumberItemT, _NumberItemT]], /
2325+
self: _nt.Array[generic[_AnyNumberItemT]],
2326+
x: _nt.Sequence1ND[_nt.op.CanRTruediv[_AnyNumberItemT, _AnyNumberItemT]],
2327+
/,
23152328
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
23162329
@overload
23172330
def __itruediv__(self: _nt.Array[object_], x: object, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
@@ -2339,7 +2352,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
23392352
) -> _nt.Array[timedelta64]: ...
23402353
@overload
23412354
def __floordiv__(
2342-
self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRFloordiv[_T]], /
2355+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanRFloordiv[_AnyItemT]], /
23432356
) -> _nt.Array[Incomplete]: ...
23442357
@overload
23452358
def __floordiv__(self: _nt.Array[object_], x: object, /) -> _nt.Array[object_]: ...
@@ -2365,7 +2378,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
23652378
def __rfloordiv__(self: _nt.Array[integer | floating], x: _nt.ToTimeDelta_nd, /) -> _nt.Array[timedelta64]: ...
23662379
@overload
23672380
def __rfloordiv__(
2368-
self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanFloordiv[_T]], /
2381+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanFloordiv[_AnyItemT]], /
23692382
) -> _nt.Array[Incomplete]: ...
23702383
@overload
23712384
def __rfloordiv__(self: _nt.Array[object_], x: object, /) -> _nt.Array[object_]: ...
@@ -2385,7 +2398,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
23852398
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
23862399
@overload
23872400
def __ifloordiv__(
2388-
self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRFloordiv[_T, _T]], /
2401+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanRFloordiv[_AnyItemT, _AnyItemT]], /
23892402
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
23902403
@overload
23912404
def __ifloordiv__(self: _nt.Array[object_], x: object, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
@@ -2408,7 +2421,9 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
24082421
@overload
24092422
def __mod__(self: _nt.Array[timedelta64], x: _nt.ToTimeDelta_nd, /) -> _nt.Array[timedelta64]: ...
24102423
@overload
2411-
def __mod__(self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRMod[_T]], /) -> _nt.Array[Incomplete]: ...
2424+
def __mod__(
2425+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanRMod[_AnyItemT]], /
2426+
) -> _nt.Array[Incomplete]: ...
24122427
@overload
24132428
def __mod__(self: _nt.Array[object_], x: object, /) -> _nt.Array[object_]: ...
24142429

@@ -2430,7 +2445,9 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
24302445
@overload
24312446
def __rmod__(self: _nt.Array[timedelta64], x: _nt.ToTimeDelta_nd, /) -> _nt.Array[timedelta64]: ...
24322447
@overload
2433-
def __rmod__(self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanMod[_T]], /) -> _nt.Array[Incomplete]: ...
2448+
def __rmod__(
2449+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanMod[_AnyItemT]], /
2450+
) -> _nt.Array[Incomplete]: ...
24342451
@overload
24352452
def __rmod__(self: _nt.Array[object_], x: object, /) -> _nt.Array[object_]: ...
24362453

@@ -2445,7 +2462,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
24452462
def __imod__(self: _nt.Array[timedelta64], x: _nt.ToTimeDelta_nd, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
24462463
@overload
24472464
def __imod__(
2448-
self: _nt.Array[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRMod[_T, _T]], /
2465+
self: _nt.Array[generic[_AnyItemT]], x: _nt.Sequence1ND[_nt.op.CanRMod[_AnyItemT, _AnyItemT]], /
24492466
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
24502467
@overload
24512468
def __imod__(self: _nt.Array[object_], x: object, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...

0 commit comments

Comments
 (0)