fix(series): arithmetics for Series[Any]

pandas-stubs/core/series.pyi

@@ -260,7 +260,8 @@ class _LocIndexerSeries(_LocIndexer, Generic[S1]):
         value: S1 | ArrayLike | Series[S1] | None,
     ) -> None: ...
 
-_ListLike: TypeAlias = (
+_ListLike: TypeAlias = ArrayLike | dict[_str, np.ndarray] | SequenceNotStr[S1]
+_ListLikeS1: TypeAlias = (
     ArrayLike | dict[_str, np.ndarray] | Sequence[S1] | IndexOpsMixin[S1]
 )
 _NumListLike: TypeAlias = (
@@ -428,7 +429,9 @@ class Series(IndexOpsMixin[S1], NDFrame):
     @overload
     def __new__(
         cls,
-        data: S1 | _ListLike[S1] | dict[HashableT1, S1] | KeysView[S1] | ValuesView[S1],
+        data: (
+            S1 | _ListLikeS1[S1] | dict[HashableT1, S1] | KeysView[S1] | ValuesView[S1]
+        ),
         index: AxesData | None = ...,
         dtype: Dtype = ...,
         name: Hashable = ...,
@@ -1628,7 +1631,9 @@ class Series(IndexOpsMixin[S1], NDFrame):
     # just failed to generate these so I couldn't match
     # them up.
     @overload
-    def __add__(self: Series[Never], other: Scalar | _ListLike | Series) -> Series: ...
+    def __add__(self: Series[Never], other: _str) -> Never: ...
+    @overload
+    def __add__(self: Series[Never], other: complex | _ListLike | Series) -> Series: ...
     @overload
     def __add__(self, other: Series[Never]) -> Series: ...
     @overload
@@ -1706,7 +1711,15 @@ class Series(IndexOpsMixin[S1], NDFrame):
     @overload
     def add(
         self: Series[Never],
-        other: Scalar | _ListLike | Series,
+        other: _str,
+        level: Level | None = None,
+        fill_value: float | None = None,
+        axis: int = 0,
+    ) -> Never: ...
+    @overload
+    def add(
+        self: Series[Never],
+        other: complex | _ListLike | Series,
         level: Level | None = None,
         fill_value: float | None = None,
         axis: int = 0,
@@ -1849,7 +1862,11 @@ class Series(IndexOpsMixin[S1], NDFrame):
         axis: int = 0,
     ) -> Series[_str]: ...
     @overload  # type: ignore[override]
-    def __radd__(self: Series[Never], other: Scalar | _ListLike) -> Series: ...
+    def __radd__(self: Series[Never], other: _str) -> Never: ...
+    @overload
+    def __radd__(
+        self: Series[Never], other: complex | _ListLike | Series
+    ) -> Series: ...
     @overload
     def __radd__(
         self: Series[bool],
@@ -1921,7 +1938,15 @@ class Series(IndexOpsMixin[S1], NDFrame):
     @overload
     def radd(
         self: Series[Never],
-        other: Scalar | _ListLike | Series,
+        other: _str,
+        level: Level | None = None,
+        fill_value: float | None = None,
+        axis: int = 0,
+    ) -> Never: ...
+    @overload
+    def radd(
+        self: Series[Never],
+        other: complex | _ListLike | Series,
         level: Level | None = None,
         fill_value: float | None = None,
         axis: int = 0,
@@ -2060,7 +2085,9 @@ class Series(IndexOpsMixin[S1], NDFrame):
         self, other: S1 | _ListLike | Series[S1] | datetime | timedelta | date
     ) -> Series[_bool]: ...
     @overload
-    def __mul__(self: Series[Never], other: complex | _ListLike | Series) -> Series: ...
+    def __mul__(
+        self: Series[Never], other: complex | _NumListLike | Series
+    ) -> Series: ...
     @overload
     def __mul__(self, other: Series[Never]) -> Series: ...  # type: ignore[overload-overlap]
     @overload
@@ -2255,7 +2282,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
     ) -> TimedeltaSeries: ...
     @overload
     def __rmul__(
-        self: Series[Never], other: complex | _ListLike | Series
+        self: Series[Never], other: complex | _NumListLike | Series
     ) -> Series: ...
     @overload
     def __rmul__(self, other: Series[Never]) -> Series: ...  # type: ignore[overload-overlap]
@@ -2892,8 +2919,8 @@ class Series(IndexOpsMixin[S1], NDFrame):
         axis: int = 0,
     ) -> Series[complex]: ...
     @overload
-    def __truediv__(
-        self: Series[Never], other: complex | _ListLike | Series
+    def __truediv__(  # type:ignore[overload-overlap]
+        self: Series[Never], other: complex | _NumListLike | Series
     ) -> Series: ...
     @overload
     def __truediv__(self, other: Series[Never]) -> Series: ...
@@ -3088,8 +3115,8 @@ class Series(IndexOpsMixin[S1], NDFrame):
     ) -> Series: ...
     div = truediv
     @overload
-    def __rtruediv__(
-        self: Series[Never], other: complex | _ListLike | Series
+    def __rtruediv__(  # type:ignore[overload-overlap]
+        self: Series[Never], other: complex | _NumListLike | Series
     ) -> Series: ...
     @overload
     def __rtruediv__(self, other: Series[Never]) -> Series: ...

tests/series/arithmetic/str/test_add.py

@@ -4,8 +4,9 @@
 import numpy as np
 from numpy import typing as npt  # noqa: F401
 import pandas as pd
+import pytest
 from typing_extensions import (
-    Never,
+    assert_never,
     assert_type,
 )
 
@@ -72,7 +73,8 @@ def test_add_numpy_array() -> None:
     r0 = np.array(["a", "bc", "def"], np.str_)
 
     if TYPE_CHECKING_INVALID_USAGE:
-        assert_type(left + i, Never)
+        with pytest.raises(AssertionError):
+            assert_never(left + i)
     check(assert_type(left + r0, "pd.Series[str]"), pd.Series, str)
 
     # `numpy` typing gives `npt.NDArray[np.int64]` in the static type

tests/series/arithmetic/test_add.py

@@ -1,9 +1,16 @@
 import numpy as np
 from numpy import typing as npt  # noqa: F401
 import pandas as pd
-from typing_extensions import assert_type
+import pytest
+from typing_extensions import (
+    assert_never,
+    assert_type,
+)
 
-from tests import check
+from tests import (
+    TYPE_CHECKING_INVALID_USAGE,
+    check,
+)
 
 # left operands
 left_i = pd.DataFrame({"a": [1, 2, 3]})["a"]
@@ -131,3 +138,18 @@ def test_add_i_pd_series() -> None:
     check(assert_type(left_i.radd(i), pd.Series), pd.Series)
     check(assert_type(left_i.radd(f), pd.Series), pd.Series)
     check(assert_type(left_i.radd(c), pd.Series), pd.Series)
+
+
+def test_add_str_py_str() -> None:
+    """Test pd.Series[Any] (int) + Python str"""
+    s = "abc"
+
+    if TYPE_CHECKING_INVALID_USAGE:
+        with pytest.raises(AssertionError):
+            assert_never(left_i + s)
+        with pytest.raises(AssertionError):
+            assert_never(s + left_i)
+        with pytest.raises(AssertionError):
+            assert_never(left_i.add(s))
+        with pytest.raises(AssertionError):
+            assert_never(left_i.radd(s))

tests/series/arithmetic/test_mul.py

@@ -3,7 +3,10 @@
 import pandas as pd
 from typing_extensions import assert_type
 
-from tests import check
+from tests import (
+    TYPE_CHECKING_INVALID_USAGE,
+    check,
+)
 
 left_i = pd.DataFrame({"a": [1, 2, 3]})["a"]  # left operand
 
@@ -129,3 +132,14 @@ def test_mul_pd_series() -> None:
     check(assert_type(left_i.rmul(i), pd.Series), pd.Series)
     check(assert_type(left_i.rmul(f), pd.Series), pd.Series)
     check(assert_type(left_i.rmul(c), pd.Series), pd.Series)
+
+
+def test_mul_str_py_str() -> None:
+    """Test pd.Series[Any] (int) * Python str"""
+    s = "abc"
+
+    if TYPE_CHECKING_INVALID_USAGE:
+        left_i * s  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
+        s * left_i  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
+        left_i.mul(s)  # type: ignore[type-var] # pyright: ignore[reportArgumentType,reportCallIssue]
+        left_i.rmul(s)  # type: ignore[type-var] # pyright: ignore[reportArgumentType,reportCallIssue]

tests/series/arithmetic/test_sub.py

@@ -7,8 +7,10 @@
 import numpy as np
 from numpy import typing as npt  # noqa: F401
 import pandas as pd
+import pytest
 from typing_extensions import (
     Never,
+    assert_never,
     assert_type,
 )
 
@@ -179,8 +181,6 @@ def test_sub_ts_numpy_datetime() -> None:
     a = np.array([s + np.timedelta64(m, "m") for m in range(3)], np.datetime64)
 
     if TYPE_CHECKING_INVALID_USAGE:
-        # `numpy` typing gives the corresponding `ndarray`s in the static type
-        # checking, where our `__rsub__` cannot override.
         _0 = left_ts - s  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
         # _1 = left_ts - a
         _2 = left_td - s  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
@@ -210,35 +210,39 @@ def test_sub_ts_pd_datetime() -> None:
     if TYPE_CHECKING_INVALID_USAGE:
         _0 = left_ts - s  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
         assert_type(left_ts - a, Never)
+
         _2 = left_td - s  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
         assert_type(left_td - a, Never)
 
         _4 = s - left_ts  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
         assert_type(a - left_ts, Never)
+
         _6 = s - left_td  # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
         assert_type(a - left_td, Never)
 
         left_ts.sub(s)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
-        assert_type(left_ts.sub(a), Never)
+        with pytest.raises(AssertionError):
+            assert_never(left_ts.sub(a))
 
+        left_td.sub(s)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
+        with pytest.raises(AssertionError):
+            assert_never(left_td.sub(a))
 
-def test_sub_ts_pd_datetime_1() -> None:
-    """Test pd.Series[Any] (Timestamp | Timedelta) - Pandas datetime(s)"""
-    s = pd.Timestamp(anchor)
-    a = pd.Series([s + pd.Timedelta(minutes=m) for m in range(3)])
-
-    if TYPE_CHECKING_INVALID_USAGE:
-        # When I merge this one to the previous one, mypy does not allow me to pass
         left_ts.rsub(s)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
-        assert_type(left_ts.rsub(a), Never)
+        with pytest.raises(AssertionError):
+            assert_never(left_ts.rsub(a))
 
+        left_td.rsub(s)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
+        with pytest.raises(AssertionError):
+            assert_never(left_td.rsub(a))
 
-def test_sub_ts_pd_datetime_2() -> None:
-    """Test pd.Series[Any] (Timestamp | Timedelta) - Pandas datetime(s)"""
-    s = pd.Timestamp(anchor)
-    a = pd.Series([s + pd.Timedelta(minutes=m) for m in range(3)])
+
+def test_sub_str_py_str() -> None:
+    """Test pd.Series[Any] (int) - Python str"""
+    s = "abc"
 
     if TYPE_CHECKING_INVALID_USAGE:
-        # When I merge this one to the previous one, mypy does not allow me to pass
-        left_td.rsub(s)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
-        assert_type(left_td.rsub(a), Never)
+        _0 = left_i - s  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
+        _1 = s - left_i  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
+        left_i.sub(s)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
+        left_i.rsub(s)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]

tests/series/arithmetic/test_truediv.py

@@ -7,6 +7,7 @@
 
 from tests import (
     PD_LTE_23,
+    TYPE_CHECKING_INVALID_USAGE,
     check,
 )
 
@@ -211,3 +212,18 @@ def test_truediv_path(tmp_path: Path) -> None:
 
         check(assert_type(fnames.rtruediv(tmp_path), pd.Series), pd.Series, Path)
         check(assert_type(fnames.rdiv(tmp_path), pd.Series), pd.Series, Path)
+
+
+def test_truediv_str_py_str() -> None:
+    """Test pd.Series[Any] (int) / Python str"""
+    s = "abc"
+
+    if TYPE_CHECKING_INVALID_USAGE:
+        _0 = left_i / s  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
+        _1 = s / left_i  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
+
+        left_i.truediv(s)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
+        left_i.div(s)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
+
+        left_i.rtruediv(s)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
+        left_i.rdiv(s)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]

tests/test_frame.py

@@ -74,6 +74,7 @@
 
 if TYPE_CHECKING:
     from pandas.core.frame import _PandasNamedTuple
+    from pandas.core.series import TimestampSeries
 else:
     _PandasNamedTuple: TypeAlias = tuple
 
@@ -4440,7 +4441,7 @@ def test_frame_setitem_na() -> None:
     df.loc[ind, :] = pd.NA
     df.iloc[[0, 2], :] = pd.NA
 
-    df["x"] = df["y"] + pd.Timedelta(days=3)
+    df["x"] = cast("TimestampSeries", df["y"]) + pd.Timedelta(days=3)
     df.loc[ind, :] = pd.NaT
     df.iloc[[0, 2], :] = pd.NaT