Skip to content

Commit 8d2a17f

Browse files
committed
feat: mul for timedelta indexes
1 parent 0444358 commit 8d2a17f

File tree

10 files changed

+340
-33
lines changed

10 files changed

+340
-33
lines changed

pandas-stubs/core/indexes/base.pyi

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -763,6 +763,22 @@ class Index(IndexOpsMixin[S1]):
763763
other: timedelta | Sequence[Timedelta] | np.timedelta64 | np_ndarray_td,
764764
) -> TimedeltaIndex: ...
765765
@overload
766+
def __mul__(self: Index[Timedelta], other: np_ndarray_complex) -> Never: ...
767+
@overload
768+
def __mul__(
769+
self: Index[Timedelta],
770+
other: (
771+
float
772+
| Sequence[float]
773+
| np_ndarray_bool
774+
| np_ndarray_anyint
775+
| np_ndarray_float
776+
| Index[bool]
777+
| Index[int]
778+
| Index[float]
779+
),
780+
) -> Index[Timedelta]: ...
781+
@overload
766782
def __mul__(self: Index[T_INT], other: bool | Sequence[bool]) -> Index[T_INT]: ...
767783
@overload
768784
def __mul__(self: Index[float], other: int | Sequence[int]) -> Index[float]: ...
@@ -818,6 +834,22 @@ class Index(IndexOpsMixin[S1]):
818834
other: timedelta | Sequence[Timedelta] | np.timedelta64 | np_ndarray_td,
819835
) -> TimedeltaIndex: ...
820836
@overload
837+
def __rmul__(self: Index[Timedelta], other: np_ndarray_complex) -> Never: ...
838+
@overload
839+
def __rmul__(
840+
self: Index[Timedelta],
841+
other: (
842+
float
843+
| Sequence[float]
844+
| np_ndarray_bool
845+
| np_ndarray_anyint
846+
| np_ndarray_float
847+
| Index[bool]
848+
| Index[int]
849+
| Index[float]
850+
),
851+
) -> Index[Timedelta]: ...
852+
@overload
821853
def __rmul__(self: Index[T_INT], other: bool | Sequence[bool]) -> Index[T_INT]: ...
822854
@overload
823855
def __rmul__(self: Index[float], other: int | Sequence[int]) -> Index[float]: ...

pandas-stubs/core/indexes/timedeltas.pyi

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,10 @@ from pandas.core.indexes.datetimelike import DatetimeTimedeltaMixin
2020
from pandas.core.indexes.datetimes import DatetimeIndex
2121
from pandas.core.indexes.period import PeriodIndex
2222
from pandas.core.series import Series
23-
from typing_extensions import Self
23+
from typing_extensions import (
24+
Never,
25+
Self,
26+
)
2427

2528
from pandas._libs import Timedelta
2629
from pandas._libs.tslibs import BaseOffset
@@ -29,6 +32,7 @@ from pandas._typing import (
2932
TimedeltaConvertibleTypes,
3033
np_ndarray_anyint,
3134
np_ndarray_bool,
35+
np_ndarray_complex,
3236
np_ndarray_dt,
3337
np_ndarray_float,
3438
np_ndarray_td,
@@ -78,9 +82,12 @@ class TimedeltaIndex(
7882
def __rsub__( # pyright: ignore[reportIncompatibleMethodOverride]
7983
self, other: dt.datetime | np.datetime64 | np_ndarray_dt | DatetimeIndex
8084
) -> DatetimeIndex: ...
81-
def __mul__( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
85+
@overload # type: ignore[override]
86+
def __mul__(self, other: np_ndarray_complex) -> Never: ...
87+
@overload
88+
def __mul__(
8289
self,
83-
other: ( # type: ignore[override]
90+
other: (
8491
float
8592
| Sequence[float]
8693
| np_ndarray_bool
@@ -91,9 +98,12 @@ class TimedeltaIndex(
9198
| Index[float]
9299
),
93100
) -> Self: ...
94-
def __rmul__( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
101+
@overload # type: ignore[override]
102+
def __rmul__(self, other: np_ndarray_complex) -> Never: ...
103+
@overload
104+
def __rmul__(
95105
self,
96-
other: ( # type: ignore[override]
106+
other: (
97107
float
98108
| Sequence[float]
99109
| np_ndarray_bool

tests/indexes/arithmetic/float/test_mul.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import numpy as np
88
from numpy import typing as npt # noqa: F401
99
import pandas as pd
10+
import pytest
1011
from typing_extensions import (
1112
Never,
1213
assert_type,
@@ -17,10 +18,15 @@
1718
check,
1819
)
1920

20-
left = pd.Index([1.0, 2.0, 3.0]) # left operand
2121

22+
@pytest.fixture
23+
def left() -> "pd.Index[float]":
24+
"""left operand"""
25+
lo = pd.Index([1.0, 2.0, 3.0])
26+
return check(assert_type(lo, "pd.Index[float]"), pd.Index, np.floating)
2227

23-
def test_mul_py_scalar() -> None:
28+
29+
def test_mul_py_scalar(left: "pd.Index[float]") -> None:
2430
"""Test pd.Index[float] * Python native scalars"""
2531
b, i, f, c = True, 1, 1.0, 1j
2632
s, d = datetime(2025, 9, 30), timedelta(seconds=1)
@@ -42,7 +48,7 @@ def test_mul_py_scalar() -> None:
4248
check(assert_type(d * left, "pd.TimedeltaIndex"), pd.TimedeltaIndex, timedelta)
4349

4450

45-
def test_mul_py_sequence() -> None:
51+
def test_mul_py_sequence(left: "pd.Index[float]") -> None:
4652
"""Test pd.Index[float] * Python native sequences"""
4753
b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j]
4854
s = [datetime(2025, 9, d) for d in (27, 28, 29)]
@@ -65,7 +71,7 @@ def test_mul_py_sequence() -> None:
6571
check(assert_type(d * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
6672

6773

68-
def test_mul_numpy_array() -> None:
74+
def test_mul_numpy_array(left: "pd.Index[float]") -> None:
6975
"""Test pd.Index[float] * numpy arrays"""
7076
b = np.array([True, False, True], np.bool_)
7177
i = np.array([2, 3, 5], np.int64)
@@ -102,7 +108,7 @@ def test_mul_numpy_array() -> None:
102108
)
103109

104110

105-
def test_mul_pd_index() -> None:
111+
def test_mul_pd_index(left: "pd.Index[float]") -> None:
106112
"""Test pd.Index[float] * pandas Indexes"""
107113
b = pd.Index([True, False, True])
108114
i = pd.Index([2, 3, 5])

tests/indexes/arithmetic/int/test_mul.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import numpy as np
88
from numpy import typing as npt # noqa: F401
99
import pandas as pd
10+
import pytest
1011
from typing_extensions import (
1112
Never,
1213
assert_type,
@@ -17,10 +18,15 @@
1718
check,
1819
)
1920

20-
left = pd.Index([1, 2, 3]) # left operand
2121

22+
@pytest.fixture
23+
def left() -> "pd.Index[int]":
24+
"""left operand"""
25+
lo = pd.Index([1, 2, 3])
26+
return check(assert_type(lo, "pd.Index[int]"), pd.Index, np.integer)
2227

23-
def test_mul_py_scalar() -> None:
28+
29+
def test_mul_py_scalar(left: "pd.Index[int]") -> None:
2430
"""Test pd.Index[int] * Python native scalars"""
2531
b, i, f, c = True, 1, 1.0, 1j
2632
s, d = datetime(2025, 9, 27), timedelta(seconds=1)
@@ -42,7 +48,7 @@ def test_mul_py_scalar() -> None:
4248
check(assert_type(d * left, pd.TimedeltaIndex), pd.Index, pd.Timedelta)
4349

4450

45-
def test_mul_py_sequence() -> None:
51+
def test_mul_py_sequence(left: "pd.Index[int]") -> None:
4652
"""Test pd.Index[int] * Python native sequences"""
4753
b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j]
4854
s = [datetime(2025, 9, d) for d in (27, 28, 29)]
@@ -67,7 +73,7 @@ def test_mul_py_sequence() -> None:
6773
check(assert_type(d * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
6874

6975

70-
def test_mul_numpy_array() -> None:
76+
def test_mul_numpy_array(left: "pd.Index[int]") -> None:
7177
"""Test pd.Index[int] * numpy arrays"""
7278
b = np.array([True, False, True], np.bool_)
7379
i = np.array([2, 3, 5], np.int64)
@@ -100,7 +106,7 @@ def test_mul_numpy_array() -> None:
100106
check(assert_type(d * left, "npt.NDArray[np.timedelta64]"), pd.Index, pd.Timedelta)
101107

102108

103-
def test_mul_pd_index() -> None:
109+
def test_mul_pd_index(left: "pd.Index[int]") -> None:
104110
"""Test pd.Index[int] * pandas Indexes"""
105111
b = pd.Index([True, False, True])
106112
i = pd.Index([2, 3, 5])
Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
from datetime import timedelta
2+
3+
import numpy as np
4+
from numpy import typing as npt # noqa: F401
5+
import pandas as pd
6+
import pytest
7+
from typing_extensions import (
8+
Never,
9+
assert_type,
10+
)
11+
12+
from tests import (
13+
PD_LTE_23,
14+
TYPE_CHECKING_INVALID_USAGE,
15+
check,
16+
)
17+
18+
19+
@pytest.fixture
20+
def left() -> "pd.Index[pd.Timedelta]":
21+
"""left operand"""
22+
# pandas-dev/pandas#62524
23+
lo = pd.Index([1]) * [timedelta(seconds=1)] # left operand
24+
return check(assert_type(lo, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
25+
26+
27+
def test_mul_py_scalar(left: "pd.Index[pd.Timedelta]") -> None:
28+
"""Test pd.Series[pd.Timedelta] * Python native scalars"""
29+
b, i, f, c = True, 1, 1.0, 1j
30+
31+
# pandas-dev/pandas#62316
32+
if PD_LTE_23:
33+
check(assert_type(left * b, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
34+
check(assert_type(left * i, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
35+
check(assert_type(left * f, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
36+
if TYPE_CHECKING_INVALID_USAGE:
37+
_0 = left * c # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
38+
39+
if PD_LTE_23:
40+
check(assert_type(b * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
41+
check(assert_type(i * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
42+
check(assert_type(f * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
43+
if TYPE_CHECKING_INVALID_USAGE:
44+
_1 = c * left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
45+
46+
47+
def test_mul_py_sequence(left: "pd.Index[pd.Timedelta]") -> None:
48+
"""Test pd.Series[pd.Timedelta] * Python native sequences"""
49+
b, i, f, c = [True], [2], [1.5], [1.7j]
50+
51+
# pandas-dev/pandas#62316
52+
if PD_LTE_23:
53+
check(assert_type(left * b, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
54+
check(assert_type(left * i, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
55+
check(assert_type(left * f, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
56+
if TYPE_CHECKING_INVALID_USAGE:
57+
_0 = left * c # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
58+
59+
if PD_LTE_23:
60+
check(assert_type(b * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
61+
check(assert_type(i * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
62+
check(assert_type(f * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
63+
if TYPE_CHECKING_INVALID_USAGE:
64+
_1 = c * left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
65+
66+
67+
def test_mul_numpy_array(left: "pd.Index[pd.Timedelta]") -> None:
68+
"""Test pd.Series[pd.Timedelta] * numpy arrays"""
69+
b = np.array([True], np.bool_)
70+
i = np.array([2], np.int64)
71+
f = np.array([1.5], np.float64)
72+
c = np.array([1.7j], np.complex128)
73+
74+
# pandas-dev/pandas#62316
75+
if PD_LTE_23:
76+
check(assert_type(left * b, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
77+
check(assert_type(left * i, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
78+
check(assert_type(left * f, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
79+
if TYPE_CHECKING_INVALID_USAGE:
80+
assert_type(left * c, Never)
81+
82+
# `numpy` typing gives the corresponding `ndarray`s in the static type
83+
# checking, where our `__rmul__` cannot override. At runtime, they return
84+
# `Series` with the correct element type.
85+
if PD_LTE_23:
86+
check(assert_type(b * left, "npt.NDArray[np.bool_]"), pd.Index, timedelta)
87+
check(assert_type(i * left, "npt.NDArray[np.int64]"), pd.Index, timedelta)
88+
check(assert_type(f * left, "npt.NDArray[np.float64]"), pd.Index, timedelta)
89+
if TYPE_CHECKING_INVALID_USAGE:
90+
# We made it Never, but numpy takes over
91+
assert_type(c * left, "npt.NDArray[np.complex128]")
92+
93+
94+
def test_mul_pd_index(left: "pd.Index[pd.Timedelta]") -> None:
95+
"""Test pd.Series[pd.Timedelta] * pandas Indexes"""
96+
b = pd.Index([True])
97+
i = pd.Index([2])
98+
f = pd.Index([1.5])
99+
c = pd.Index([1.7j])
100+
101+
# pandas-dev/pandas#62316
102+
if PD_LTE_23:
103+
check(assert_type(left * b, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
104+
check(assert_type(left * i, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
105+
check(assert_type(left * f, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
106+
if TYPE_CHECKING_INVALID_USAGE:
107+
_0 = left * c # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
108+
109+
if PD_LTE_23:
110+
check(assert_type(b * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
111+
check(assert_type(i * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
112+
check(assert_type(f * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
113+
if TYPE_CHECKING_INVALID_USAGE:
114+
_1 = c * left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]

tests/indexes/timedeltaindex/__init__.py

Whitespace-only changes.

0 commit comments

Comments
 (0)