diff --git a/pandas-stubs/core/base.pyi b/pandas-stubs/core/base.pyi index e9447a76f..55af19d4f 100644 --- a/pandas-stubs/core/base.pyi +++ b/pandas-stubs/core/base.pyi @@ -29,6 +29,7 @@ from pandas._libs.tslibs.timedeltas import Timedelta from pandas._typing import ( S1, S2, + S3, AxisIndex, DropKeep, DTypeLike, @@ -322,6 +323,30 @@ class ElementOpsMixin(Generic[S2]): def _proto_rfloordiv( self: ElementOpsMixin[Timedelta], other: timedelta | np.timedelta64 | Timedelta ) -> ElementOpsMixin[int]: ... + @overload + def _proto_mod( + self: ElementOpsMixin[int], other: int | np.integer + ) -> ElementOpsMixin[int]: ... + @overload + def _proto_mod( + self: ElementOpsMixin[float], other: float | np.floating + ) -> ElementOpsMixin[float]: ... + @overload + def _proto_mod( + self: ElementOpsMixin[Timedelta], other: timedelta | np.timedelta64 | Timedelta + ) -> ElementOpsMixin[Timedelta]: ... + @overload + def _proto_rmod( + self: ElementOpsMixin[int], other: int | np.integer + ) -> ElementOpsMixin[int]: ... + @overload + def _proto_rmod( + self: ElementOpsMixin[float], other: float | np.floating + ) -> ElementOpsMixin[float]: ... + @overload + def _proto_rmod( + self: ElementOpsMixin[Timedelta], other: timedelta | np.timedelta64 | Timedelta + ) -> ElementOpsMixin[Timedelta]: ... @type_check_only class Supports_ProtoAdd(Protocol[_T_contra, S2]): @@ -354,3 +379,21 @@ class Supports_ProtoFloorDiv(Protocol[_T_contra, S2]): @type_check_only class Supports_ProtoRFloorDiv(Protocol[_T_contra, S2]): def _proto_rfloordiv(self, other: _T_contra, /) -> ElementOpsMixin[S2]: ... + +@type_check_only +class Supports_ProtoMod(Protocol[_T_contra, S2]): + def _proto_mod(self, other: _T_contra, /) -> ElementOpsMixin[S2]: ... + +@type_check_only +class Supports_ProtoRMod(Protocol[_T_contra, S2]): + def _proto_rmod(self, other: _T_contra, /) -> ElementOpsMixin[S2]: ... + +@type_check_only +class Supports_ProtoDivMod(Protocol[_T_contra, S2, S3]): + def _proto_floordiv(self, other: _T_contra, /) -> ElementOpsMixin[S2]: ... + def _proto_mod(self, other: _T_contra, /) -> ElementOpsMixin[S3]: ... + +@type_check_only +class Supports_ProtoRDivMod(Protocol[_T_contra, S2, S3]): + def _proto_rfloordiv(self, other: _T_contra, /) -> ElementOpsMixin[S2]: ... + def _proto_rmod(self, other: _T_contra, /) -> ElementOpsMixin[S3]: ... diff --git a/tests/test_natype.py b/tests/test_natype.py index 7a524bc37..c55d07a47 100644 --- a/tests/test_natype.py +++ b/tests/test_natype.py @@ -83,11 +83,8 @@ def test_arithmetic() -> None: # __divmod__ # bug upstream: https://github.com/pandas-dev/pandas/issues/62196 # check( - # assert_type( - # divmod(na, s_int), - # tuple[pd.Series, pd.Series], - # ), - # tuple, + assert_type(divmod(na, s_int), tuple[pd.Series, pd.Series]) + # , tuple # ) # check( # assert_type( @@ -110,13 +107,13 @@ def test_arithmetic() -> None: # __rdivmod__ # bug upstream: https://github.com/pandas-dev/pandas/issues/62196 # check( - # assert_type(divmod(s_int, na), tuple[pd.Series, pd.Series]), - # tuple, + assert_type(divmod(s_int, na), "tuple[pd.Series[int], pd.Series[int]]") + # , tuple, # ) - # https://github.com/pandas-dev/pandas-stubs/issues/1347 + # bug upstream: https://github.com/pandas-dev/pandas/issues/62196 # check( - # assert_type(divmod(idx_int, na), tuple[pd.Index, pd.Index]), - # tuple, + assert_type(divmod(idx_int, na), "tuple[pd.Index[int], pd.Index[int]]") + # , tuple, # ) check(assert_type(divmod(1, na), tuple[NAType, NAType]), tuple)