Skip to content

Commit 5cba4c1

Browse files
committed
refactor: truediv
1 parent 9105941 commit 5cba4c1

File tree

9 files changed

+386
-285
lines changed

9 files changed

+386
-285
lines changed

pandas-stubs/_libs/tslibs/timedeltas.pyi

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -255,22 +255,24 @@ class Timedelta(timedelta):
255255
) -> np_ndarray[ShapeT, np.int_]: ...
256256
# Override due to more types supported than timedelta
257257
@overload # type: ignore[override]
258-
def __truediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
258+
# pyrefly: ignore[bad-override]
259+
def __truediv__(self, other: Just[int] | Just[float]) -> Self: ...
259260
@overload
260-
def __truediv__(self, other: float) -> Self: ...
261+
def __truediv__(self, other: Self | NaTType) -> float: ...
261262
@overload
262263
def __truediv__(
263264
self, other: np_ndarray[ShapeT, np.integer | np.floating]
264265
) -> np_ndarray[ShapeT, np.timedelta64]: ...
265266
@overload
266-
def __truediv__(self, other: Series[Timedelta]) -> Series[float]: ...
267-
@overload
268-
def __truediv__(self, other: Series[int]) -> Series[Timedelta]: ...
267+
def __truediv__(
268+
self, other: np_ndarray[ShapeT, np.timedelta64]
269+
) -> np_ndarray[ShapeT, np.floating]: ...
269270
@overload
270-
def __truediv__(self, other: Series[float]) -> Series[Timedelta]: ...
271+
def __rtruediv__(self, other: Self | NaTType) -> float: ...
271272
@overload
272-
def __truediv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
273-
def __rtruediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
273+
def __rtruediv__(
274+
self, other: np_ndarray[ShapeT, np.timedelta64]
275+
) -> np_ndarray[ShapeT, np.floating]: ...
274276
# Override due to more types supported than timedelta
275277
@overload
276278
def __eq__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]

pandas-stubs/_typing.pyi

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ from typing import (
2525
overload,
2626
)
2727

28+
from _typeshed import _T_contra
2829
import numpy as np
2930
from numpy import typing as npt
3031
import pandas as pd
@@ -860,6 +861,9 @@ np_ndarray_td: TypeAlias = npt.NDArray[np.timedelta64]
860861
# Define shape and generic type variables with defaults similar to numpy
861862
GenericT = TypeVar("GenericT", bound=np.generic, default=Any)
862863
GenericT_co = TypeVar("GenericT_co", bound=np.generic, default=Any, covariant=True)
864+
GenericT_contra = TypeVar(
865+
"GenericT_contra", bound=np.generic, default=Any, contravariant=True
866+
)
863867
ShapeT = TypeVar("ShapeT", bound=tuple[int, ...], default=tuple[Any, ...])
864868
# Numpy ndarray with more ergonomic typevar
865869
np_ndarray: TypeAlias = np.ndarray[ShapeT, np.dtype[GenericT]]
@@ -1114,4 +1118,10 @@ class Just(Protocol, Generic[T]):
11141118
@override
11151119
def __class__(self, t: type[T], /) -> None: ...
11161120

1121+
class SupportsTrueDiv(Protocol[_T_contra, _T_co]):
1122+
def __truediv__(self, x: _T_contra, /) -> _T_co: ...
1123+
1124+
class SupportsRTrueDiv(Protocol[_T_contra, _T_co]):
1125+
def __rtruediv__(self, x: _T_contra, /) -> _T_co: ...
1126+
11171127
__all__ = ["npt", "type_t"]

pandas-stubs/core/base.pyi

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ from collections.abc import (
33
Iterator,
44
Sequence,
55
)
6+
from datetime import timedelta
67
from typing import (
78
Any,
89
Generic,
@@ -269,6 +270,38 @@ class ElementOpsMixin(Generic[S2]):
269270
def _proto_rmul(
270271
self: ElementOpsMixin[str], other: Just[int] | np.integer
271272
) -> ElementOpsMixin[str]: ...
273+
@overload
274+
def _proto_truediv(
275+
self: ElementOpsMixin[int], other: int | np.integer
276+
) -> ElementOpsMixin[float]: ...
277+
@overload
278+
def _proto_truediv(
279+
self: ElementOpsMixin[float], other: float | np.floating
280+
) -> ElementOpsMixin[float]: ...
281+
@overload
282+
def _proto_truediv(
283+
self: ElementOpsMixin[complex], other: complex | np.complexfloating
284+
) -> ElementOpsMixin[complex]: ...
285+
@overload
286+
def _proto_truediv(
287+
self: ElementOpsMixin[Timedelta], other: timedelta | Timedelta | np.timedelta64
288+
) -> ElementOpsMixin[float]: ...
289+
@overload
290+
def _proto_rtruediv(
291+
self: ElementOpsMixin[int], other: int | np.integer
292+
) -> ElementOpsMixin[float]: ...
293+
@overload
294+
def _proto_rtruediv(
295+
self: ElementOpsMixin[float], other: float | np.floating
296+
) -> ElementOpsMixin[float]: ...
297+
@overload
298+
def _proto_rtruediv(
299+
self: ElementOpsMixin[complex], other: complex | np.complexfloating
300+
) -> ElementOpsMixin[complex]: ...
301+
@overload
302+
def _proto_rtruediv(
303+
self: ElementOpsMixin[Timedelta], other: timedelta | Timedelta | np.timedelta64
304+
) -> ElementOpsMixin[float]: ...
272305

273306
@type_check_only
274307
class Supports_ProtoAdd(Protocol[_T_contra, S2]):
@@ -285,3 +318,11 @@ class Supports_ProtoMul(Protocol[_T_contra, S2]):
285318
@type_check_only
286319
class Supports_ProtoRMul(Protocol[_T_contra, S2]):
287320
def _proto_rmul(self, other: _T_contra, /) -> ElementOpsMixin[S2]: ...
321+
322+
@type_check_only
323+
class Supports_ProtoTrueDiv(Protocol[_T_contra, S2]):
324+
def _proto_truediv(self, other: _T_contra, /) -> ElementOpsMixin[S2]: ...
325+
326+
@type_check_only
327+
class Supports_ProtoRTrueDiv(Protocol[_T_contra, S2]):
328+
def _proto_rtruediv(self, other: _T_contra, /) -> ElementOpsMixin[S2]: ...

0 commit comments

Comments
 (0)