GH1089 Add more tests

Author: loicdiridollou

pandas-stubs/core/series.pyi

@@ -1822,15 +1822,23 @@ class Series(IndexOpsMixin[S1], NDFrame):
     # Met    @overload
     @overload
     def add(
+        self,
+        other: S1 | Self,
+        level: Level | None = ...,
+        fill_value: float | None = ...,
+        axis: int = ...,
+    ) -> Self: ...
+    @overload
+    def add(  # pyright: ignore[reportOverlappingOverload]
         self: Series[int],
-        other: int,
+        other: Series[int] | int,
         level: Level | None = ...,
         fill_value: float | None = ...,
         axis: int = ...,
     ) -> Series[int]: ...
     @overload
-    def add(  # pyright: ignore[reportOverlappingOverload]
-        self,
+    def add(
+        self: Series[int] | Series[float],
         other: float | Series[float],
         level: Level | None = ...,
         fill_value: float | None = ...,
@@ -1846,7 +1854,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
     ) -> Series[float]: ...
     @overload
     def add(
-        self,
+        self: Series[complex],
         other: complex,
         level: Level | None = ...,
         fill_value: float | None = ...,
@@ -1860,14 +1868,6 @@ class Series(IndexOpsMixin[S1], NDFrame):
         fill_value: float | None = ...,
         axis: int = ...,
     ) -> Series[S1]: ...
-    @overload
-    def add(
-        self,
-        other: S1 | Self,
-        level: Level | None = ...,
-        fill_value: float | None = ...,
-        axis: int = ...,
-    ) -> Series: ...
     def all(
         self,
         axis: AxisIndex = ...,
@@ -2279,6 +2279,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
         **kwargs,
     ) -> float: ...
     @overload
+    def sub(
+        self,
+        other: S1 | Self,
+        level: Level | None = ...,
+        fill_value: float | None = ...,
+        axis: int = ...,
+    ) -> Self: ...
+    @overload
     def sub(  # pyright: ignore[reportOverlappingOverload]
         self: Series[int],
         other: int,

tests/test_series_arithmetic.py

@@ -1,7 +1,5 @@
 """Test module for arithmetic operations on Series."""
 
-from typing import cast
-
 import numpy as np
 import pandas as pd
 from typing_extensions import assert_type
@@ -193,11 +191,12 @@ def test_element_wise_float_float() -> None:
 
 
 def test_element_wise_int_unknown() -> None:
-    s = cast(pd.Series, pd.Series([7, -5, 10]))
+    df = pd.DataFrame({"a": [7, -5, 10]})
+    s = df["a"]
     s2 = pd.Series([0, 1, -105])
 
     check(assert_type(s + s2, pd.Series), pd.Series)
-    check(assert_type(s.add(s2, fill_value=0), "pd.Series[float]"), pd.Series)
+    check(assert_type(s.add(s2, fill_value=0), pd.Series), pd.Series)
 
     check(assert_type(s - s2, pd.Series), pd.Series)
     check(assert_type(s.sub(s2, fill_value=0), pd.Series), pd.Series)
@@ -208,3 +207,75 @@ def test_element_wise_int_unknown() -> None:
     # GH1089 should be the following
     check(assert_type(s / s2, "pd.Series[float]"), pd.Series)
     check(assert_type(s.div(s2, fill_value=0), "pd.Series[float]"), pd.Series)
+
+
+def test_element_wise_unknown_int() -> None:
+    df = pd.DataFrame({"a": [7, -5, 10]})
+    s = pd.Series([0, 1, -105])
+    s2 = df["a"]
+
+    check(assert_type(s + s2, pd.Series), pd.Series)
+    check(assert_type(s.add(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s - s2, pd.Series), pd.Series)
+    check(assert_type(s.sub(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s * s2, pd.Series), pd.Series)
+    check(assert_type(s.mul(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s / s2, "pd.Series[float]"), pd.Series)
+    check(assert_type(s.div(s2, fill_value=0), "pd.Series[float]"), pd.Series)
+
+
+def test_element_wise_unknown_unknown() -> None:
+    df = pd.DataFrame({"a": [7, -5, 10]})
+    s = df["a"]
+    s2 = df["a"]
+
+    check(assert_type(s + s2, pd.Series), pd.Series)
+    check(assert_type(s.add(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s - s2, pd.Series), pd.Series)
+    check(assert_type(s.sub(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s * s2, pd.Series), pd.Series)
+    check(assert_type(s.mul(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s / s2, "pd.Series[float]"), pd.Series)
+    check(assert_type(s.div(s2, fill_value=0), "pd.Series[float]"), pd.Series)
+
+
+def test_element_wise_float_unknown() -> None:
+    df = pd.DataFrame({"a": [7, -5, 10]})
+    s = pd.Series([1.3, 2.5, 4.5])
+    s2 = df["a"]
+
+    check(assert_type(s + s2, pd.Series), pd.Series)
+    check(assert_type(s.add(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s - s2, pd.Series), pd.Series)
+    check(assert_type(s.sub(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s * s2, pd.Series), pd.Series)
+    check(assert_type(s.mul(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s / s2, "pd.Series[float]"), pd.Series)
+    check(assert_type(s.div(s2, fill_value=0), "pd.Series[float]"), pd.Series)
+
+
+def test_element_wise_unknown_float() -> None:
+    df = pd.DataFrame({"a": [7, -5, 10]})
+    s = df["a"]
+    s2 = pd.Series([1.3, 2.5, 4.5])
+
+    check(assert_type(s + s2, pd.Series), pd.Series)
+    check(assert_type(s.add(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s - s2, pd.Series), pd.Series)
+    check(assert_type(s.sub(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s * s2, pd.Series), pd.Series)
+    check(assert_type(s.mul(s2, fill_value=0), pd.Series), pd.Series)
+
+    check(assert_type(s / s2, "pd.Series[float]"), pd.Series)
+    check(assert_type(s.div(s2, fill_value=0), "pd.Series[float]"), pd.Series)