Skip to content

Commit 04f4bec

Browse files
jorenhamjorenham
committed
✅ fix failing ndarray binop tests for __floordiv__ and __mod__ 
1 parent 9f573b7 commit 04f4bec

File tree

2 files changed

+80
-111
lines changed

2 files changed

+80
-111
lines changed

‎src/numpy-stubs/__init__.pyi

Lines changed: 79 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -649,6 +649,7 @@ _ScalarT = TypeVar("_ScalarT", bound=generic)
649649
_SelfScalarT = TypeVar("_SelfScalarT", bound=generic)
650650
_ScalarT_co = TypeVar("_ScalarT_co", bound=generic, default=Any, covariant=True)
651651
_IntegralScalarT = TypeVar("_IntegralScalarT", bound=bool_ | integer | object_)
652+
_RealT = TypeVar("_RealT", bound=bool_ | integer | floating)
652653
_RealScalarT = TypeVar("_RealScalarT", bound=bool_ | integer | floating | object_)
653654
_RealNumberT = TypeVar("_RealNumberT", bound=integer | floating)
654655
_IntegerT = TypeVar("_IntegerT", bound=integer)
@@ -2541,169 +2542,137 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
25412542
@overload
25422543
def __itruediv__(self: NDArray[object_], x: object, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
25432544

2544-
# TODO(jorenham): rewrite the modular and bitwise binops below
2545-
2546-
# the pyright error appears to be a false positive
2547-
@overload
2548-
def __floordiv__(self: NDArray[_RealNumberT], rhs: int | bool_, /) -> _nt.Array[_RealNumberT, _ShapeT_co]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
2549-
@overload
2550-
def __floordiv__(self: NDArray[_RealNumberT], rhs: _nt.ToBool_nd, /) -> NDArray[_RealNumberT]: ... # type: ignore[overload-overlap]
2551-
@overload
2552-
def __floordiv__(self: NDArray[bool_], rhs: _ArrayLike[_RealNumberT], /) -> NDArray[_RealNumberT]: ...
2553-
@overload
2554-
def __floordiv__(self: NDArray[bool_], rhs: _nt.ToBool_nd, /) -> NDArray[int8]: ...
2555-
@overload
2556-
def __floordiv__(self: NDArray[float64], rhs: _nt.CoFloat64_nd, /) -> NDArray[float64]: ...
2557-
@overload
2558-
def __floordiv__(self: _ArrayFloat64_co, rhs: _nt.ToFloat64_nd, /) -> NDArray[float64]: ...
2559-
@overload
2560-
def __floordiv__(self: NDArray[unsignedinteger], rhs: _nt.CoUInt64_nd, /) -> NDArray[unsignedinteger]: ...
2561-
@overload
2562-
def __floordiv__(self: _ArrayUInt_co, rhs: _nt.ToUInteger_nd, /) -> NDArray[unsignedinteger]: ...
2563-
@overload
2564-
def __floordiv__(self: NDArray[signedinteger], rhs: _nt.CoInt64_nd, /) -> NDArray[signedinteger]: ...
2545+
#
25652546
@overload
2566-
def __floordiv__(self: _ArrayInt_co, rhs: _nt.ToSInteger_nd, /) -> NDArray[signedinteger]: ...
2547+
def __floordiv__(self: NDArray[bool_], x: _nt.ToBool_nd, /) -> NDArray[int8]: ...
25672548
@overload
2568-
def __floordiv__(self: NDArray[floating], rhs: _nt.CoFloating_nd, /) -> NDArray[floating]: ...
2549+
def __floordiv__(self: NDArray[_RealNumberT], x: _nt.Casts[_RealNumberT], /) -> NDArray[_RealNumberT]: ...
25692550
@overload
2570-
def __floordiv__(self: NDArray[floating | integer], rhs: _nt.ToFloating_nd, /) -> NDArray[floating]: ...
2551+
def __floordiv__(
2552+
self: NDArray[_RealNumberT], x: _nt.CastsWith[_RealNumberT, _RealScalarT], /
2553+
) -> NDArray[_RealScalarT]: ...
25712554
@overload
2572-
def __floordiv__(self: _nt.CoFloating_nd, rhs: _nt.CoFloating_nd, /) -> NDArray[Incomplete]: ...
2555+
def __floordiv__(self: _nt.CastsWithInt[_RealScalarT], x: _nt.SequenceND[int], /) -> NDArray[_RealScalarT]: ...
25732556
@overload
2574-
def __floordiv__(self: NDArray[timedelta64], rhs: _ArrayLike[timedelta64], /) -> NDArray[int64]: ...
2557+
def __floordiv__(
2558+
self: _nt.CastsWithFloat[_RealScalarT], x: _nt.SequenceND[_nt.JustFloat], /
2559+
) -> NDArray[_RealScalarT]: ...
25752560
@overload
2576-
def __floordiv__(self: NDArray[timedelta64], rhs: _nt.ToBool_nd, /) -> Never: ...
2561+
def __floordiv__(self: NDArray[timedelta64], x: _nt.ToTimeDelta_nd, /) -> NDArray[int64]: ...
25772562
@overload
2578-
def __floordiv__(self: NDArray[timedelta64], rhs: _nt.CoFloating_nd, /) -> NDArray[timedelta64]: ...
2563+
def __floordiv__(
2564+
self: NDArray[timedelta64], x: _nt.ToInteger_nd | _nt.ToFloating_nd, /
2565+
) -> NDArray[timedelta64]: ...
25792566
@overload
2580-
def __floordiv__(self: NDArray[object_], rhs: object, /) -> NDArray[object_]: ...
2567+
def __floordiv__(
2568+
self: NDArray[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRFloordiv[_T]], /
2569+
) -> NDArray[Incomplete]: ...
25812570
@overload
2582-
def __floordiv__(self, rhs: _ArrayLikeObject_co, /) -> NDArray[object_]: ...
2571+
def __floordiv__(self: NDArray[object_], x: object, /) -> NDArray[object_]: ...
25832572

25842573
#
25852574
@overload
2586-
def __rfloordiv__(self: NDArray[_RealNumberT], lhs: int | bool_, /) -> _nt.Array[_RealNumberT, _ShapeT_co]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
2575+
def __rfloordiv__(self: NDArray[bool_], x: _nt.ToBool_nd, /) -> NDArray[int8]: ...
25872576
@overload
2588-
def __rfloordiv__(self: NDArray[_RealNumberT], lhs: _nt.ToBool_nd, /) -> NDArray[_RealNumberT]: ... # type: ignore[overload-overlap]
2577+
def __rfloordiv__(self: NDArray[_RealNumberT], x: _nt.Casts[_RealNumberT], /) -> NDArray[_RealNumberT]: ...
25892578
@overload
2590-
def __rfloordiv__(self: NDArray[bool_], lhs: _ArrayLike[_RealNumberT], /) -> NDArray[_RealNumberT]: ...
2579+
def __rfloordiv__(
2580+
self: NDArray[_RealNumberT], x: _nt.CastsWith[_RealNumberT, _RealScalarT], /
2581+
) -> NDArray[_RealScalarT]: ...
25912582
@overload
2592-
def __rfloordiv__(self: NDArray[bool_], lhs: _nt.ToBool_nd, /) -> NDArray[int8]: ...
2583+
def __rfloordiv__(self: _nt.CastsWithInt[_RealScalarT], x: _nt.SequenceND[int], /) -> NDArray[_RealScalarT]: ...
25932584
@overload
2594-
def __rfloordiv__(self: NDArray[float64], lhs: _nt.CoFloat64_nd, /) -> NDArray[float64]: ...
2585+
def __rfloordiv__(
2586+
self: _nt.CastsWithFloat[_RealScalarT], x: _nt.SequenceND[_nt.JustFloat], /
2587+
) -> NDArray[_RealScalarT]: ...
25952588
@overload
2596-
def __rfloordiv__(self: _ArrayFloat64_co, lhs: _nt.ToFloat64_nd, /) -> NDArray[float64]: ...
2589+
def __rfloordiv__(self: NDArray[timedelta64], x: _nt.ToTimeDelta_nd, /) -> NDArray[int64]: ...
25972590
@overload
2598-
def __rfloordiv__(self: NDArray[unsignedinteger], lhs: _nt.CoUInt64_nd, /) -> NDArray[unsignedinteger]: ...
2591+
def __rfloordiv__(self: NDArray[integer | floating], x: _nt.ToTimeDelta_nd, /) -> NDArray[timedelta64]: ...
25992592
@overload
2600-
def __rfloordiv__(self: _ArrayUInt_co, lhs: _nt.ToUInteger_nd, /) -> NDArray[unsignedinteger]: ...
2601-
@overload
2602-
def __rfloordiv__(self: NDArray[signedinteger], lhs: _nt.CoInt64_nd, /) -> NDArray[signedinteger]: ...
2603-
@overload
2604-
def __rfloordiv__(self: _ArrayInt_co, lhs: _nt.ToSInteger_nd, /) -> NDArray[signedinteger]: ...
2605-
@overload
2606-
def __rfloordiv__(self: NDArray[floating], lhs: _nt.CoFloating_nd, /) -> NDArray[floating]: ...
2607-
@overload
2608-
def __rfloordiv__(self: NDArray[floating | integer], lhs: _nt.ToFloating_nd, /) -> NDArray[floating]: ...
2609-
@overload
2610-
def __rfloordiv__(self: _nt.CoFloating_nd, lhs: _nt.CoFloating_nd, /) -> NDArray[Incomplete]: ...
2611-
@overload
2612-
def __rfloordiv__(self: NDArray[timedelta64], lhs: _ArrayLike[timedelta64], /) -> NDArray[int64]: ...
2613-
@overload
2614-
def __rfloordiv__(self: NDArray[floating | integer], lhs: _ArrayLike[timedelta64], /) -> NDArray[timedelta64]: ...
2615-
@overload
2616-
def __rfloordiv__(self: NDArray[object_], lhs: object, /) -> NDArray[object_]: ...
2593+
def __rfloordiv__(
2594+
self: NDArray[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanFloordiv[_T]], /
2595+
) -> NDArray[Incomplete]: ...
26172596
@overload
2618-
def __rfloordiv__(self, lhs: _ArrayLikeObject_co, /) -> NDArray[object_]: ...
2597+
def __rfloordiv__(self: NDArray[object_], x: object, /) -> NDArray[object_]: ...
26192598

26202599
#
26212600
@overload # type: ignore[misc]
26222601
def __ifloordiv__(
2623-
self: NDArray[integer | timedelta64],
2624-
rhs: _nt.CoInteger_nd,
2625-
/,
2602+
self: NDArray[_RealNumberT], x: _nt.Casts[_RealNumberT], /
26262603
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
26272604
@overload
2628-
def __ifloordiv__(self: NDArray[floating], rhs: _nt.CoFloating_nd, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
2605+
def __ifloordiv__(self: NDArray[integer], x: _nt.SequenceND[int], /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
26292606
@overload
2630-
def __ifloordiv__(self: NDArray[complexfloating], rhs: _nt.CoComplex_nd, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
2631-
@overload
2632-
def __ifloordiv__(self: NDArray[object_], rhs: object, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
2633-
2634-
#
2607+
def __ifloordiv__(self: NDArray[floating], x: _nt.SequenceND[float], /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
26352608
@overload
2636-
def __mod__(self: NDArray[_RealNumberT], rhs: int | bool_, /) -> _nt.Array[_RealNumberT, _ShapeT_co]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
2637-
@overload
2638-
def __mod__(self: NDArray[_RealNumberT], rhs: _nt.ToBool_nd, /) -> NDArray[_RealNumberT]: ... # type: ignore[overload-overlap]
2639-
@overload
2640-
def __mod__(self: NDArray[bool_], rhs: _ArrayLike[_RealNumberT], /) -> NDArray[_RealNumberT]: ...
2641-
@overload
2642-
def __mod__(self: NDArray[bool_], rhs: _nt.ToBool_nd, /) -> NDArray[int8]: ...
2643-
@overload
2644-
def __mod__(self: NDArray[float64], rhs: _nt.CoFloat64_nd, /) -> NDArray[float64]: ...
2609+
def __ifloordiv__(
2610+
self: NDArray[timedelta64], x: _nt.ToInteger_nd | _nt.ToFloating_nd, /
2611+
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
26452612
@overload
2646-
def __mod__(self: _ArrayFloat64_co, rhs: _nt.ToFloat64_nd, /) -> NDArray[float64]: ...
2613+
def __ifloordiv__(
2614+
self: NDArray[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRFloordiv[_T, _T]], /
2615+
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
26472616
@overload
2648-
def __mod__(self: NDArray[unsignedinteger], rhs: _nt.CoUInt64_nd, /) -> NDArray[unsignedinteger]: ...
2617+
def __ifloordiv__(self: NDArray[object_], x: object, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
2618+
2619+
#
26492620
@overload
2650-
def __mod__(self: _ArrayUInt_co, rhs: _nt.ToUInteger_nd, /) -> NDArray[unsignedinteger]: ...
2621+
def __mod__(self: NDArray[bool_], x: _nt.ToBool_nd, /) -> NDArray[int8]: ...
26512622
@overload
2652-
def __mod__(self: NDArray[signedinteger], rhs: _nt.CoInt64_nd, /) -> NDArray[signedinteger]: ...
2623+
def __mod__(self: NDArray[_RealNumberT], x: _nt.Casts[_RealNumberT], /) -> NDArray[_RealNumberT]: ...
26532624
@overload
2654-
def __mod__(self: _ArrayInt_co, rhs: _nt.ToSInteger_nd, /) -> NDArray[signedinteger]: ...
2625+
def __mod__(self: NDArray[_RealT], x: _nt.CastsWith[_RealT, _RealScalarT], /) -> NDArray[_RealScalarT]: ...
26552626
@overload
2656-
def __mod__(self: NDArray[floating], rhs: _nt.CoFloating_nd, /) -> NDArray[floating]: ...
2627+
def __mod__(self: _nt.CastsWithInt[_RealScalarT], x: _nt.SequenceND[int], /) -> NDArray[_RealScalarT]: ...
26572628
@overload
2658-
def __mod__(self: NDArray[floating | integer], rhs: _nt.ToFloating_nd, /) -> NDArray[floating]: ...
2629+
def __mod__(
2630+
self: _nt.CastsWithFloat[_RealScalarT], x: _nt.SequenceND[_nt.JustFloat], /
2631+
) -> NDArray[_RealScalarT]: ...
26592632
@overload
2660-
def __mod__(self: NDArray[timedelta64], rhs: _ArrayLike[timedelta64], /) -> NDArray[timedelta64]: ...
2633+
def __mod__(self: NDArray[timedelta64], x: _nt.ToTimeDelta_nd, /) -> NDArray[timedelta64]: ...
26612634
@overload
2662-
def __mod__(self: NDArray[object_], rhs: object, /) -> NDArray[object_]: ...
2635+
def __mod__(self: NDArray[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRMod[_T]], /) -> NDArray[Incomplete]: ...
26632636
@overload
2664-
def __mod__(self, rhs: _ArrayLikeObject_co, /) -> NDArray[object_]: ...
2637+
def __mod__(self: NDArray[object_], x: object, /) -> NDArray[object_]: ...
26652638

2666-
# keep in sync with __mod__
2667-
@overload
2668-
def __rmod__(self: NDArray[_RealNumberT], lhs: int | bool_, /) -> _nt.Array[_RealNumberT, _ShapeT_co]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
2669-
@overload
2670-
def __rmod__(self: NDArray[_RealNumberT], lhs: _nt.ToBool_nd, /) -> NDArray[_RealNumberT]: ... # type: ignore[overload-overlap]
2671-
@overload
2672-
def __rmod__(self: NDArray[bool_], lhs: _ArrayLike[_RealNumberT], /) -> NDArray[_RealNumberT]: ...
2673-
@overload
2674-
def __rmod__(self: NDArray[bool_], lhs: _nt.ToBool_nd, /) -> NDArray[int8]: ...
2675-
@overload
2676-
def __rmod__(self: NDArray[float64], lhs: _nt.CoFloat64_nd, /) -> NDArray[float64]: ...
2677-
@overload
2678-
def __rmod__(self: _ArrayFloat64_co, lhs: _nt.ToFloat64_nd, /) -> NDArray[float64]: ...
2679-
@overload
2680-
def __rmod__(self: NDArray[unsignedinteger], lhs: _nt.CoUInt64_nd, /) -> NDArray[unsignedinteger]: ...
2639+
#
26812640
@overload
2682-
def __rmod__(self: _ArrayUInt_co, lhs: _nt.ToUInteger_nd, /) -> NDArray[unsignedinteger]: ...
2641+
def __rmod__(self: NDArray[bool_], x: _nt.ToBool_nd, /) -> NDArray[int8]: ...
26832642
@overload
2684-
def __rmod__(self: NDArray[signedinteger], lhs: _nt.CoInt64_nd, /) -> NDArray[signedinteger]: ...
2643+
def __rmod__(self: NDArray[_RealNumberT], x: _nt.Casts[_RealNumberT], /) -> NDArray[_RealNumberT]: ...
26852644
@overload
2686-
def __rmod__(self: _ArrayInt_co, lhs: _nt.ToSInteger_nd, /) -> NDArray[signedinteger]: ...
2645+
def __rmod__(self: NDArray[_RealT], x: _nt.CastsWith[_RealT, _RealScalarT], /) -> NDArray[_RealScalarT]: ...
26872646
@overload
2688-
def __rmod__(self: NDArray[floating], lhs: _nt.CoFloating_nd, /) -> NDArray[floating]: ...
2647+
def __rmod__(self: _nt.CastsWithInt[_RealScalarT], x: _nt.SequenceND[int], /) -> NDArray[_RealScalarT]: ...
26892648
@overload
2690-
def __rmod__(self: NDArray[floating | integer], lhs: _nt.ToFloating_nd, /) -> NDArray[floating]: ...
2649+
def __rmod__(
2650+
self: _nt.CastsWithFloat[_RealScalarT], x: _nt.SequenceND[_nt.JustFloat], /
2651+
) -> NDArray[_RealScalarT]: ...
26912652
@overload
2692-
def __rmod__(self: NDArray[timedelta64], lhs: _ArrayLike[timedelta64], /) -> NDArray[timedelta64]: ...
2653+
def __rmod__(self: NDArray[timedelta64], x: _nt.ToTimeDelta_nd, /) -> NDArray[timedelta64]: ...
26932654
@overload
2694-
def __rmod__(self: NDArray[object_], lhs: object, /) -> NDArray[object_]: ...
2655+
def __rmod__(self: NDArray[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanMod[_T]], /) -> NDArray[Incomplete]: ...
26952656
@overload
2696-
def __rmod__(self, lhs: _ArrayLikeObject_co, /) -> NDArray[object_]: ...
2657+
def __rmod__(self: NDArray[object_], x: object, /) -> NDArray[object_]: ...
26972658

26982659
#
26992660
@overload # type: ignore[misc]
2700-
def __imod__(self: NDArray[integer], rhs: _nt.CoInteger_nd, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
2661+
def __imod__(self: NDArray[_RealNumberT], x: _nt.Casts[_RealNumberT], /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
2662+
@overload
2663+
def __imod__(self: NDArray[integer], x: _nt.SequenceND[int], /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
27012664
@overload
2702-
def __imod__(self: NDArray[floating], rhs: _nt.CoFloating_nd, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
2665+
def __imod__(self: NDArray[floating], x: _nt.SequenceND[float], /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
27032666
@overload
2704-
def __imod__(self: NDArray[timedelta64], rhs: _ArrayLike[timedelta64], /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
2667+
def __imod__(self: NDArray[timedelta64], x: _nt.ToTimeDelta_nd, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
27052668
@overload
2706-
def __imod__(self: NDArray[object_], rhs: object, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
2669+
def __imod__(
2670+
self: NDArray[generic[_T]], x: _nt.Sequence1ND[_nt.op.CanRMod[_T, _T]], /
2671+
) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
2672+
@overload
2673+
def __imod__(self: NDArray[object_], x: object, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
2674+
2675+
# TODO(jorenham): rewrite the binops below
27072676

27082677
#
27092678
@overload

‎src/numpy-stubs/ma/core.pyi

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -465,7 +465,7 @@ class MaskedArray(np.ndarray[_ShapeT_co, _DTypeT_co]):
465465
def __imul__(self, other: Incomplete, /) -> Self: ... # type: ignore[override]
466466
def __idiv__(self, other: Incomplete, /) -> Self: ...
467467
@override
468-
def __ifloordiv__(self, other: Incomplete, /) -> Self: ... # type: ignore[override]
468+
def __ifloordiv__(self, other: Incomplete, /) -> Self: ...
469469
@override
470470
def __itruediv__(self, other: Incomplete, /) -> Self: ...
471471
@override

0 commit comments

Comments
 (0)