feat: mul for timedelta indexes

Author: cmp0xff

pandas-stubs/core/indexes/base.pyi

@@ -751,6 +751,22 @@ class Index(IndexOpsMixin[S1]):
         other: timedelta | Sequence[Timedelta] | np.timedelta64 | np_ndarray_td,
     ) -> TimedeltaIndex: ...
     @overload
+    def __mul__(self: Index[Timedelta], other: np_ndarray_complex) -> Never: ...
+    @overload
+    def __mul__(
+        self: Index[Timedelta],
+        other: (
+            float
+            | Sequence[float]
+            | np_ndarray_bool
+            | np_ndarray_anyint
+            | np_ndarray_float
+            | Index[bool]
+            | Index[int]
+            | Index[float]
+        ),
+    ) -> Index[Timedelta]: ...
+    @overload
     def __mul__(self: Index[T_INT], other: bool | Sequence[bool]) -> Index[T_INT]: ...
     @overload
     def __mul__(self: Index[float], other: int | Sequence[int]) -> Index[float]: ...
@@ -806,6 +822,22 @@ class Index(IndexOpsMixin[S1]):
         other: timedelta | Sequence[Timedelta] | np.timedelta64 | np_ndarray_td,
     ) -> TimedeltaIndex: ...
     @overload
+    def __rmul__(self: Index[Timedelta], other: np_ndarray_complex) -> Never: ...
+    @overload
+    def __rmul__(
+        self: Index[Timedelta],
+        other: (
+            float
+            | Sequence[float]
+            | np_ndarray_bool
+            | np_ndarray_anyint
+            | np_ndarray_float
+            | Index[bool]
+            | Index[int]
+            | Index[float]
+        ),
+    ) -> Index[Timedelta]: ...
+    @overload
     def __rmul__(self: Index[T_INT], other: bool | Sequence[bool]) -> Index[T_INT]: ...
     @overload
     def __rmul__(self: Index[float], other: int | Sequence[int]) -> Index[float]: ...

pandas-stubs/core/indexes/timedeltas.pyi

@@ -20,7 +20,10 @@ from pandas.core.indexes.datetimelike import DatetimeTimedeltaMixin
 from pandas.core.indexes.datetimes import DatetimeIndex
 from pandas.core.indexes.period import PeriodIndex
 from pandas.core.series import Series
-from typing_extensions import Self
+from typing_extensions import (
+    Never,
+    Self,
+)
 
 from pandas._libs import Timedelta
 from pandas._libs.tslibs import BaseOffset
@@ -29,6 +32,7 @@ from pandas._typing import (
     TimedeltaConvertibleTypes,
     np_ndarray_anyint,
     np_ndarray_bool,
+    np_ndarray_complex,
     np_ndarray_dt,
     np_ndarray_float,
     np_ndarray_td,
@@ -78,9 +82,12 @@ class TimedeltaIndex(
     def __rsub__(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, other: dt.datetime | np.datetime64 | np_ndarray_dt | DatetimeIndex
     ) -> DatetimeIndex: ...
-    def __mul__(  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
+    @overload  # type: ignore[override]
+    def __mul__(self, other: np_ndarray_complex) -> Never: ...
+    @overload
+    def __mul__(
         self,
-        other: (  # type: ignore[override]
+        other: (
             float
             | Sequence[float]
             | np_ndarray_bool
@@ -91,9 +98,12 @@ class TimedeltaIndex(
             | Index[float]
         ),
     ) -> Self: ...
-    def __rmul__(  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
+    @overload  # type: ignore[override]
+    def __rmul__(self, other: np_ndarray_complex) -> Never: ...
+    @overload
+    def __rmul__(
         self,
-        other: (  # type: ignore[override]
+        other: (
             float
             | Sequence[float]
             | np_ndarray_bool

pyproject.toml

@@ -42,7 +42,7 @@ pyarrow = ">=10.0.1"
 pytest = ">=8.4.2"
 pyright = ">=1.1.405"
 ty = ">=0.0.1a21"
-pyrefly = ">=0.34.0"
+pyrefly = ">=0.35.0"
 poethepoet = ">=0.16.5"
 loguru = ">=0.6.0"
 typing-extensions = ">=4.4.0"

tests/indexes/arithmetic/float/test_mul.py

@@ -7,6 +7,7 @@
 import numpy as np
 from numpy import typing as npt  # noqa: F401
 import pandas as pd
+import pytest
 from typing_extensions import (
     Never,
     assert_type,
@@ -17,10 +18,15 @@
     check,
 )
 
-left = pd.Index([1.0, 2.0, 3.0])  # left operand
 
+@pytest.fixture
+def left() -> "pd.Index[float]":
+    """left operand"""
+    lo = pd.Index([1.0, 2.0, 3.0])
+    return check(assert_type(lo, "pd.Index[float]"), pd.Index, np.floating)
 
-def test_mul_py_scalar() -> None:
+
+def test_mul_py_scalar(left: "pd.Index[float]") -> None:
     """Test pd.Index[float] * Python native scalars"""
     b, i, f, c = True, 1, 1.0, 1j
     s, d = datetime(2025, 9, 30), timedelta(seconds=1)
@@ -42,7 +48,7 @@ def test_mul_py_scalar() -> None:
     check(assert_type(d * left, "pd.TimedeltaIndex"), pd.TimedeltaIndex, timedelta)
 
 
-def test_mul_py_sequence() -> None:
+def test_mul_py_sequence(left: "pd.Index[float]") -> None:
     """Test pd.Index[float] * Python native sequences"""
     b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j]
     s = [datetime(2025, 9, d) for d in (27, 28, 29)]
@@ -65,7 +71,7 @@ def test_mul_py_sequence() -> None:
     check(assert_type(d * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
 
 
-def test_mul_numpy_array() -> None:
+def test_mul_numpy_array(left: "pd.Index[float]") -> None:
     """Test pd.Index[float] * numpy arrays"""
     b = np.array([True, False, True], np.bool_)
     i = np.array([2, 3, 5], np.int64)
@@ -102,7 +108,7 @@ def test_mul_numpy_array() -> None:
     )
 
 
-def test_mul_pd_index() -> None:
+def test_mul_pd_index(left: "pd.Index[float]") -> None:
     """Test pd.Index[float] * pandas Indexes"""
     b = pd.Index([True, False, True])
     i = pd.Index([2, 3, 5])

tests/indexes/arithmetic/int/test_mul.py

@@ -7,6 +7,7 @@
 import numpy as np
 from numpy import typing as npt  # noqa: F401
 import pandas as pd
+import pytest
 from typing_extensions import (
     Never,
     assert_type,
@@ -17,10 +18,15 @@
     check,
 )
 
-left = pd.Index([1, 2, 3])  # left operand
 
+@pytest.fixture
+def left() -> "pd.Index[int]":
+    """left operand"""
+    lo = pd.Index([1, 2, 3])
+    return check(assert_type(lo, "pd.Index[int]"), pd.Index, np.integer)
 
-def test_mul_py_scalar() -> None:
+
+def test_mul_py_scalar(left: "pd.Index[int]") -> None:
     """Test pd.Index[int] * Python native scalars"""
     b, i, f, c = True, 1, 1.0, 1j
     s, d = datetime(2025, 9, 27), timedelta(seconds=1)
@@ -42,7 +48,7 @@ def test_mul_py_scalar() -> None:
     check(assert_type(d * left, pd.TimedeltaIndex), pd.Index, pd.Timedelta)
 
 
-def test_mul_py_sequence() -> None:
+def test_mul_py_sequence(left: "pd.Index[int]") -> None:
     """Test pd.Index[int] * Python native sequences"""
     b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j]
     s = [datetime(2025, 9, d) for d in (27, 28, 29)]
@@ -67,7 +73,7 @@ def test_mul_py_sequence() -> None:
     check(assert_type(d * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
 
 
-def test_mul_numpy_array() -> None:
+def test_mul_numpy_array(left: "pd.Index[int]") -> None:
     """Test pd.Index[int] * numpy arrays"""
     b = np.array([True, False, True], np.bool_)
     i = np.array([2, 3, 5], np.int64)
@@ -100,7 +106,7 @@ def test_mul_numpy_array() -> None:
     check(assert_type(d * left, "npt.NDArray[np.timedelta64]"), pd.Index, pd.Timedelta)
 
 
-def test_mul_pd_index() -> None:
+def test_mul_pd_index(left: "pd.Index[int]") -> None:
     """Test pd.Index[int] * pandas Indexes"""
     b = pd.Index([True, False, True])
     i = pd.Index([2, 3, 5])

tests/indexes/arithmetic/timedelta/test_mul.py

@@ -0,0 +1,114 @@
+from datetime import timedelta
+
+import numpy as np
+from numpy import typing as npt  # noqa: F401
+import pandas as pd
+import pytest
+from typing_extensions import (
+    Never,
+    assert_type,
+)
+
+from tests import (
+    PD_LTE_23,
+    TYPE_CHECKING_INVALID_USAGE,
+    check,
+)
+
+
+@pytest.fixture
+def left() -> "pd.Index[pd.Timedelta]":
+    """left operand"""
+    # pandas-dev/pandas#62524
+    lo = pd.Index([1]) * [timedelta(seconds=1)]  # left operand
+    return check(assert_type(lo, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+
+
+def test_mul_py_scalar(left: "pd.Index[pd.Timedelta]") -> None:
+    """Test pd.Series[pd.Timedelta] * Python native scalars"""
+    b, i, f, c = True, 1, 1.0, 1j
+
+    # pandas-dev/pandas#62316
+    if PD_LTE_23:
+        check(assert_type(left * b, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(left * i, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(left * f, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    if TYPE_CHECKING_INVALID_USAGE:
+        _0 = left * c  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
+
+    if PD_LTE_23:
+        check(assert_type(b * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(i * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(f * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    if TYPE_CHECKING_INVALID_USAGE:
+        _1 = c * left  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
+
+
+def test_mul_py_sequence(left: "pd.Index[pd.Timedelta]") -> None:
+    """Test pd.Series[pd.Timedelta] * Python native sequences"""
+    b, i, f, c = [True], [2], [1.5], [1.7j]
+
+    # pandas-dev/pandas#62316
+    if PD_LTE_23:
+        check(assert_type(left * b, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(left * i, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(left * f, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    if TYPE_CHECKING_INVALID_USAGE:
+        _0 = left * c  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
+
+    if PD_LTE_23:
+        check(assert_type(b * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(i * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(f * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    if TYPE_CHECKING_INVALID_USAGE:
+        _1 = c * left  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
+
+
+def test_mul_numpy_array(left: "pd.Index[pd.Timedelta]") -> None:
+    """Test pd.Series[pd.Timedelta] * numpy arrays"""
+    b = np.array([True], np.bool_)
+    i = np.array([2], np.int64)
+    f = np.array([1.5], np.float64)
+    c = np.array([1.7j], np.complex128)
+
+    # pandas-dev/pandas#62316
+    if PD_LTE_23:
+        check(assert_type(left * b, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(left * i, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(left * f, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    if TYPE_CHECKING_INVALID_USAGE:
+        assert_type(left * c, Never)
+
+    # `numpy` typing gives the corresponding `ndarray`s in the static type
+    # checking, where our `__rmul__` cannot override. At runtime, they return
+    # `Series` with the correct element type.
+    if PD_LTE_23:
+        check(assert_type(b * left, "npt.NDArray[np.bool_]"), pd.Index, timedelta)
+    check(assert_type(i * left, "npt.NDArray[np.int64]"), pd.Index, timedelta)
+    check(assert_type(f * left, "npt.NDArray[np.float64]"), pd.Index, timedelta)
+    if TYPE_CHECKING_INVALID_USAGE:
+        # We made it Never, but numpy takes over
+        assert_type(c * left, "npt.NDArray[np.complex128]")
+
+
+def test_mul_pd_index(left: "pd.Index[pd.Timedelta]") -> None:
+    """Test pd.Series[pd.Timedelta] * pandas Indexes"""
+    b = pd.Index([True])
+    i = pd.Index([2])
+    f = pd.Index([1.5])
+    c = pd.Index([1.7j])
+
+    # pandas-dev/pandas#62316
+    if PD_LTE_23:
+        check(assert_type(left * b, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(left * i, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(left * f, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    if TYPE_CHECKING_INVALID_USAGE:
+        _0 = left * c  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
+
+    if PD_LTE_23:
+        check(assert_type(b * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(i * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    check(assert_type(f * left, "pd.Index[pd.Timedelta]"), pd.Index, timedelta)
+    if TYPE_CHECKING_INVALID_USAGE:
+        _1 = c * left  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]