Skip to content

Commit a22516c

Browse files
cmp0xffcmp0xff
authored
feat(arithmetic): 🪜 floordiv (#1452)
* floordiv * remove irrelevant changes * floordiv * remove irrelevant change * floordiv * pyright * mypy * pytest * pyrefly * fix: np.ndarray // pd.IndexOpsMixin * mypy * series * mypy * series * series * fix: pandas-dev/pandas#63007 pytest * py310 * index * comment * pre-commit * comments * #1474 * numpy/numpy#30173
1 parent 8a52340 commit a22516c

File tree

24 files changed

+2218
-704
lines changed

24 files changed

+2218
-704
lines changed

‎pandas-stubs/_typing.pyi‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -854,6 +854,9 @@ np_ndarray_anyint: TypeAlias = npt.NDArray[np.integer]
854854
np_ndarray_float: TypeAlias = npt.NDArray[np.floating]
855855
np_ndarray_complex: TypeAlias = npt.NDArray[np.complexfloating]
856856
np_ndarray_bool: TypeAlias = npt.NDArray[np.bool_]
857+
np_ndarray_num: TypeAlias = npt.NDArray[
858+
np.bool | np.integer | np.floating | np.complexfloating
859+
]
857860
np_ndarray_str: TypeAlias = npt.NDArray[np.str_]
858861
np_ndarray_dt: TypeAlias = npt.NDArray[np.datetime64]
859862
np_ndarray_td: TypeAlias = npt.NDArray[np.timedelta64]
@@ -1114,8 +1117,7 @@ Incomplete: TypeAlias = Any
11141117
class Just(Protocol, Generic[T]):
11151118
@property # type: ignore[override]
11161119
@override
1117-
# pyrefly: ignore # bad-override
1118-
def __class__(self, /) -> type[T]: ...
1120+
def __class__(self, /) -> type[T]: ... # pyrefly: ignore[bad-override]
11191121
@__class__.setter
11201122
@override
11211123
def __class__(self, t: type[T], /) -> None: ...

‎pandas-stubs/core/base.pyi‎

Lines changed: 116 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,11 @@ import numpy as np
2020
from pandas.core.arraylike import OpsMixin
2121
from pandas.core.arrays import ExtensionArray
2222
from pandas.core.arrays.categorical import Categorical
23+
from pandas.core.arrays.integer import IntegerArray
24+
from pandas.core.arrays.timedeltas import TimedeltaArray
2325
from pandas.core.indexes.accessors import ArrayDescriptor
2426
from pandas.core.indexes.base import Index
27+
from pandas.core.indexes.timedeltas import TimedeltaIndex
2528
from pandas.core.series import Series
2629
from typing_extensions import Self
2730

@@ -44,6 +47,7 @@ from pandas._typing import (
4447
np_ndarray_bool,
4548
np_ndarray_complex,
4649
np_ndarray_float,
50+
np_ndarray_td,
4751
)
4852
from pandas.util._decorators import cache_readonly
4953

@@ -168,7 +172,84 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
168172
) -> np.intp: ...
169173
def drop_duplicates(self, *, keep: DropKeep = ...) -> Self: ...
170174

171-
NumListLike: TypeAlias = (
175+
ScalarArrayIndexJustInt: TypeAlias = (
176+
Just[int]
177+
| np.integer
178+
| Sequence[Just[int] | np.integer]
179+
| np_ndarray_anyint
180+
| IntegerArray
181+
| Index[int]
182+
)
183+
ScalarArrayIndexSeriesJustInt: TypeAlias = ScalarArrayIndexJustInt | Series[int]
184+
ScalarArrayIndexJustFloat: TypeAlias = (
185+
Just[float]
186+
| np.floating
187+
| Sequence[Just[float] | np.floating]
188+
| np_ndarray_float
189+
# | FloatingArray # TODO: after pandas-dev/pandas-stubs#1469
190+
| Index[float]
191+
)
192+
ScalarArrayIndexSeriesJustFloat: TypeAlias = ScalarArrayIndexJustFloat | Series[float]
193+
ScalarArrayIndexJustComplex: TypeAlias = (
194+
Just[complex]
195+
| np.complexfloating
196+
| Sequence[Just[complex] | np.complexfloating]
197+
| np_ndarray_complex
198+
| Index[complex]
199+
)
200+
ScalarArrayIndexSeriesJustComplex: TypeAlias = (
201+
ScalarArrayIndexJustComplex | Series[complex]
202+
)
203+
204+
ScalarArrayIndexIntNoBool: TypeAlias = (
205+
Just[int]
206+
| np.integer
207+
| Sequence[int | np.integer]
208+
| np_ndarray_anyint
209+
| IntegerArray
210+
| Index[int]
211+
)
212+
ScalarArrayIndexSeriesIntNoBool: TypeAlias = ScalarArrayIndexIntNoBool | Series[int]
213+
214+
NumpyRealScalar: TypeAlias = np.bool | np.integer | np.floating
215+
IndexReal: TypeAlias = Index[bool] | Index[int] | Index[float]
216+
ScalarArrayIndexReal: TypeAlias = (
217+
float
218+
| Sequence[float | NumpyRealScalar]
219+
| NumpyRealScalar
220+
| np.typing.NDArray[NumpyRealScalar]
221+
| ExtensionArray
222+
| IndexReal
223+
)
224+
SeriesReal: TypeAlias = Series[bool] | Series[int] | Series[float]
225+
ScalarArrayIndexSeriesReal: TypeAlias = ScalarArrayIndexReal | SeriesReal
226+
227+
NumpyComplexScalar: TypeAlias = NumpyRealScalar | np.complexfloating
228+
IndexComplex: TypeAlias = IndexReal | Index[complex]
229+
ScalarArrayIndexComplex: TypeAlias = (
230+
complex
231+
| Sequence[complex | NumpyComplexScalar]
232+
| NumpyComplexScalar
233+
| np.typing.NDArray[NumpyComplexScalar]
234+
| ExtensionArray
235+
| IndexComplex
236+
)
237+
SeriesComplex: TypeAlias = SeriesReal | Series[complex]
238+
ScalarArrayIndexSeriesComplex: TypeAlias = ScalarArrayIndexComplex | SeriesComplex
239+
240+
ArrayIndexTimedeltaNoSeq: TypeAlias = np_ndarray_td | TimedeltaArray | TimedeltaIndex
241+
ScalarArrayIndexTimedelta: TypeAlias = (
242+
timedelta
243+
| np.timedelta64
244+
| Sequence[timedelta | np.timedelta64]
245+
| ArrayIndexTimedeltaNoSeq
246+
)
247+
ArrayIndexSeriesTimedeltaNoSeq: TypeAlias = ArrayIndexTimedeltaNoSeq | Series[Timedelta]
248+
ScalarArrayIndexSeriesTimedelta: TypeAlias = (
249+
ScalarArrayIndexTimedelta | Series[Timedelta]
250+
)
251+
252+
NumListLike: TypeAlias = ( # TODO: pandas-dev/pandas-stubs#1474 deprecated, do not use
172253
ExtensionArray
173254
| np_ndarray_bool
174255
| np_ndarray_anyint
@@ -280,7 +361,7 @@ class ElementOpsMixin(Generic[S2]):
280361
) -> ElementOpsMixin[complex]: ...
281362
@overload
282363
def _proto_truediv(
283-
self: ElementOpsMixin[Timedelta], other: timedelta | Timedelta | np.timedelta64
364+
self: ElementOpsMixin[Timedelta], other: timedelta | np.timedelta64 | Timedelta
284365
) -> ElementOpsMixin[float]: ...
285366
@overload
286367
def _proto_rtruediv(
@@ -296,8 +377,32 @@ class ElementOpsMixin(Generic[S2]):
296377
) -> ElementOpsMixin[complex]: ...
297378
@overload
298379
def _proto_rtruediv(
299-
self: ElementOpsMixin[Timedelta], other: timedelta | Timedelta | np.timedelta64
380+
self: ElementOpsMixin[Timedelta], other: timedelta | np.timedelta64 | Timedelta
381+
) -> ElementOpsMixin[float]: ...
382+
@overload
383+
def _proto_floordiv(
384+
self: ElementOpsMixin[int], other: int | np.integer
385+
) -> ElementOpsMixin[int]: ...
386+
@overload
387+
def _proto_floordiv(
388+
self: ElementOpsMixin[float], other: float | np.floating
300389
) -> ElementOpsMixin[float]: ...
390+
@overload
391+
def _proto_floordiv(
392+
self: ElementOpsMixin[Timedelta], other: timedelta | np.timedelta64 | Timedelta
393+
) -> ElementOpsMixin[int]: ...
394+
@overload
395+
def _proto_rfloordiv(
396+
self: ElementOpsMixin[int], other: int | np.integer
397+
) -> ElementOpsMixin[int]: ...
398+
@overload
399+
def _proto_rfloordiv(
400+
self: ElementOpsMixin[float], other: float | np.floating
401+
) -> ElementOpsMixin[float]: ...
402+
@overload
403+
def _proto_rfloordiv(
404+
self: ElementOpsMixin[Timedelta], other: timedelta | np.timedelta64 | Timedelta
405+
) -> ElementOpsMixin[int]: ...
301406

302407
@type_check_only
303408
class Supports_ProtoAdd(Protocol[_T_contra, S2]):
@@ -322,3 +427,11 @@ class Supports_ProtoTrueDiv(Protocol[_T_contra, S2]):
322427
@type_check_only
323428
class Supports_ProtoRTrueDiv(Protocol[_T_contra, S2]):
324429
def _proto_rtruediv(self, other: _T_contra, /) -> ElementOpsMixin[S2]: ...
430+
431+
@type_check_only
432+
class Supports_ProtoFloorDiv(Protocol[_T_contra, S2]):
433+
def _proto_floordiv(self, other: _T_contra, /) -> ElementOpsMixin[S2]: ...
434+
435+
@type_check_only
436+
class Supports_ProtoRFloorDiv(Protocol[_T_contra, S2]):
437+
def _proto_rfloordiv(self, other: _T_contra, /) -> ElementOpsMixin[S2]: ...

0 commit comments

Comments
 (0)