feat(series): implement the proposal

Author: cmp0xff

pandas-stubs/core/series.pyi

@@ -1631,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
@@ -1709,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,
@@ -1852,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],
@@ -1924,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,

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,9 +7,11 @@
 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,
+    Never
 )
 
 from tests import (
@@ -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,45 +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))
 
-
-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)])
-
-    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)
+        with pytest.raises(AssertionError):
+            assert_never(left_td.rsub(a))
+
 
+def test_sub_str_py_str() -> None:
+    """Test pd.Series[Any] (int) - Python str"""
+    s = "abc"
 
-def test_str_sub() -> None:
     if TYPE_CHECKING_INVALID_USAGE:
-        left_i - "abc"  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
-        "abc" - left_i  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
-        left_i * "abc"  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
-        "abc" * left_i  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
-        left_i / "abc"  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
-        "abc" / left_i  # type: ignore[operator] # pyright:ignore[reportOperatorIssue]
+        _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