diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 8b78c6d16..8eeb225c0 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -857,7 +857,7 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def dot(self, other: Series[S1]) -> Scalar: ... @overload - def dot(self, other: DataFrame) -> Series[S1]: ... + def dot(self, other: DataFrame) -> Series: ... @overload def dot( self, other: ArrayLike | dict[_str, np.ndarray] | Sequence[S1] | Index[S1] @@ -1609,6 +1609,21 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __add__(self, other: S1 | Self) -> Self: ... @overload + def __add__( + self: Series[int] | Series[float], + other: float | Series[float], + ) -> Series[float]: ... + @overload + def __add__( + self: Series[float], + other: int | Series[int], + ) -> Series[float]: ... + @overload + def __add__( + self, + other: complex, + ) -> Series[complex]: ... + @overload def __add__( self, other: num | _str | timedelta | Timedelta | _ListLike | Series | np.timedelta64, @@ -1623,8 +1638,10 @@ class Series(IndexOpsMixin[S1], NDFrame): self, other: int | np_ndarray_anyint | Series[int] ) -> Series[int]: ... # def __array__(self, dtype: Optional[_bool] = ...) -> _np_ndarray - def __div__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... def __eq__(self, other: object) -> Series[_bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + @overload + def __floordiv__(self: Series[float], other: Series[float] | Series[int] | float) -> Series[float]: ... # type: ignore[overload-overlap] + @overload def __floordiv__(self, other: num | _ListLike | Series[S1]) -> Series[int]: ... def __ge__( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] self, other: S1 | _ListLike | Series[S1] | datetime | timedelta | date @@ -1639,13 +1656,43 @@ class Series(IndexOpsMixin[S1], NDFrame): self, other: S1 | _ListLike | Series[S1] | datetime | timedelta | date ) -> Series[_bool]: ... @overload + def __mul__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: S1 | Self + ) -> Self: ... + @overload def __mul__( self, other: timedelta | Timedelta | TimedeltaSeries | np.timedelta64 ) -> TimedeltaSeries: ... @overload + def __mul__( # pyright: ignore[reportOverlappingOverload] + self: Series[int], other: Series[int] | int + ) -> Series[int]: ... + @overload + def __mul__(self: Series[int], other: Series[float] | float) -> Series[float]: ... + @overload + def __mul__(self: Series[float], other: Series[int] | int) -> Series[float]: ... + @overload + def __mul__(self: Series[Any], other: Series[Any]) -> Series: ... + @overload def __mul__(self, other: num | _ListLike | Series) -> Series: ... + @overload + def __mod__( # pyright: ignore[reportOverlappingOverload] + self: Series[int], other: int | Series[int] + ) -> Series[int]: ... + @overload + def __mod__(self: Series[int], other: float | Series[float]) -> Series[float]: ... + @overload + def __mod__( + self: Series[float], other: float | Series[int] | Series[float] + ) -> Series[float]: ... + @overload def __mod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... def __ne__(self, other: object) -> Series[_bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + @overload + def __pow__(self: Series[int], other: Series[int] | int) -> Series[int]: ... + @overload + def __pow__(self, other: Series[float] | float) -> Series[float]: ... + @overload def __pow__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... # ignore needed for mypy as we want different results based on the arguments @overload # type: ignore[override] @@ -1669,9 +1716,13 @@ class Series(IndexOpsMixin[S1], NDFrame): def __rand__( # pyright: ignore[reportIncompatibleMethodOverride] self, other: int | np_ndarray_anyint | Series[int] ) -> Series[int]: ... - def __rdiv__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... def __rdivmod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] - def __rfloordiv__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... + @overload + def __rfloordiv__(self: Series[float], other: float | Series[int] | Series[float]) -> Series[float]: ... # type: ignore[misc] + @overload + def __rfloordiv__( + self, other: num | _ListLike | Series[float] + ) -> Series[float]: ... def __rmod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... @overload def __rmul__( @@ -1716,8 +1767,35 @@ class Series(IndexOpsMixin[S1], NDFrame): self, other: Timestamp | datetime | TimestampSeries ) -> TimedeltaSeries: ... @overload + def __sub__(self, other: S1 | Self) -> Self: ... + @overload + def __sub__( # pyright: ignore[reportOverlappingOverload] + self: Series[int], + other: int | Series[int], + ) -> Series[int]: ... + @overload + def __sub__( + self: Series[int], + other: float | Series[float], + ) -> Series[float]: ... + @overload + def __sub__( + self: Series[float], + other: int | Series[int], + ) -> Series[float]: ... + @overload + def __sub__( + self, + other: complex, + ) -> Series[complex]: ... + @overload def __sub__(self, other: num | _ListLike | Series) -> Series: ... - def __truediv__(self, other: num | _ListLike | Series[S1] | Path) -> Series: ... + @overload + def __truediv__(self: Series[int], other: Series[int] | int) -> Series[float]: ... + @overload + def __truediv__(self, other: Path) -> Series: ... + @overload + def __truediv__(self, other: num | _ListLike | Series[S1]) -> Series[float]: ... # ignore needed for mypy as we want different results based on the arguments @overload # type: ignore[override] def __xor__( # pyright: ignore[reportOverlappingOverload] @@ -1741,7 +1819,48 @@ class Series(IndexOpsMixin[S1], NDFrame): def iloc(self) -> _iLocIndexerSeries[S1]: ... @property def loc(self) -> _LocIndexerSeries[S1]: ... - # Methods + # 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: Series[int] | int, + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[int]: ... + @overload + def add( + self: Series[int] | Series[float], + other: float | Series[float], + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[float]: ... + @overload + def add( + self: Series[float], + other: int | Series[int], + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[float]: ... + @overload + def add( + self: Series[complex], + other: complex, + level: Level | None = ..., + fill_value: float | None = ..., + axis: int = ..., + ) -> Series[complex]: ... + @overload def add( self, other: Series[S1] | Scalar, @@ -1812,13 +1931,30 @@ class Series(IndexOpsMixin[S1], NDFrame): min_periods: int = ..., method: CalculationMethod = ..., ) -> Expanding[Series]: ... + @overload + def floordiv( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self: Series[int], + other: Series[int] | int, + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[int]: ... + @overload + def floordiv( + self: Series[int] | Series[float], + other: num | _ListLike | Series[S1], + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[float]: ... + @overload def floordiv( self, other: num | _ListLike | Series[S1], level: Level | None = ..., fill_value: float | None = ..., axis: AxisIndex | None = ..., - ) -> Series[int]: ... + ) -> Self: ... def ge( self, other: Scalar | Series[S1], @@ -1896,6 +2032,31 @@ class Series(IndexOpsMixin[S1], NDFrame): numeric_only: _bool = ..., **kwargs, ) -> S1: ... + @overload + def mod( # pyright: ignore[reportOverlappingOverload] + self: Series[int], + other: Series[int] | int, + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[int]: ... + @overload + def mod( + self: Series[int], + other: Series[float] | float, + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[float]: ... + @overload + def mod( + self: Series[float], + other: Series[int] | Series[float] | float, + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[float]: ... + @overload def mod( self, other: num | _ListLike | Series[S1], @@ -1904,6 +2065,38 @@ class Series(IndexOpsMixin[S1], NDFrame): axis: AxisIndex | None = ..., ) -> Series[S1]: ... @overload + def mul( # type: ignore[overload-overlap] + self, + other: Series[S1] | Self, + level: Level | None = ..., + fill_value: int | None = ..., + axis: AxisIndex | None = ..., + ) -> Self: ... + @overload + def mul( # pyright: ignore[reportOverlappingOverload] + self: Series[int], + other: Series[int] | int, + level: Level | None = ..., + fill_value: int | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[int]: ... + @overload + def mul( + self: Series[float], + other: Series[int] | Series[float] | int, + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[float]: ... + @overload + def mul( + self: Series[int], + other: Series[float] | float, + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[float]: ... + @overload def mul( self, other: timedelta | Timedelta | TimedeltaSeries | np.timedelta64, @@ -1918,7 +2111,7 @@ class Series(IndexOpsMixin[S1], NDFrame): level: Level | None = ..., fill_value: float | None = ..., axis: AxisIndex | None = ..., - ) -> Series: ... + ) -> Self: ... def multiply( self, other: num | _ListLike | Series[S1], @@ -1934,6 +2127,15 @@ class Series(IndexOpsMixin[S1], NDFrame): axis: AxisIndex = ..., ) -> Series[_bool]: ... def nunique(self, dropna: _bool = ...) -> int: ... + @overload + def pow( + self: Series[int], + other: float | Series[float], + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[float]: ... + @overload def pow( self, other: num | _ListLike | Series[S1], @@ -2076,6 +2278,39 @@ class Series(IndexOpsMixin[S1], NDFrame): numeric_only: _bool = ..., **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, + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[int]: ... + @overload + def sub( + self: Series[int], + other: float | Series[float], + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[float]: ... + @overload + def sub( + self, + other: complex, + level: Level | None = ..., + fill_value: float | None = ..., + axis: AxisIndex | None = ..., + ) -> Series[complex]: ... + @overload def sub( self, other: num | _ListLike | Series[S1], diff --git a/tests/test_frame.py b/tests/test_frame.py index ad7bb6450..f9015985b 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2451,10 +2451,9 @@ def test_sum_get_add() -> None: def test_getset_untyped() -> None: - result: int = 10 df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]}) # Tests that Dataframe.__getitem__ needs to return untyped series. - result = df["x"].max() + check(assert_type(df["x"].max(), Any), np.integer) def test_getmultiindex_columns() -> None: @@ -2966,6 +2965,7 @@ def sum_mean(x: pd.DataFrame) -> float: ) lfunc: Callable[[pd.DataFrame], float] = lambda x: x.sum().mean() + with pytest_warns_bounded( DeprecationWarning, "DataFrameGroupBy.apply operated on the grouping columns.", diff --git a/tests/test_series.py b/tests/test_series.py index f0dbe8019..36966318c 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -145,20 +145,20 @@ def test_types_all() -> None: def test_types_csv() -> None: s = pd.Series(data=[1, 2, 3]) - csv_df: str = s.to_csv() + check(assert_type(s.to_csv(), str), str) with ensure_clean() as path: s.to_csv(path) - s2: pd.DataFrame = pd.read_csv(path) + check(assert_type(pd.read_csv(path), pd.DataFrame), pd.DataFrame) with ensure_clean() as path: s.to_csv(Path(path)) - s3: pd.DataFrame = pd.read_csv(Path(path)) + check(assert_type(pd.read_csv(Path(path)), pd.DataFrame), pd.DataFrame) # This keyword was added in 1.1.0 https://pandas.pydata.org/docs/whatsnew/v1.1.0.html with ensure_clean() as path: s.to_csv(path, errors="replace") - s4: pd.DataFrame = pd.read_csv(path) + check(assert_type(pd.read_csv(path), pd.DataFrame), pd.DataFrame) def test_types_copy() -> None: @@ -229,11 +229,11 @@ def test_types_boolean_indexing() -> None: def test_types_df_to_df_comparison() -> None: s = pd.Series(data={"col1": [1, 2]}) s2 = pd.Series(data={"col1": [3, 2]}) - res_gt: pd.Series = s > s2 - res_ge: pd.Series = s >= s2 - res_lt: pd.Series = s < s2 - res_le: pd.Series = s <= s2 - res_e: pd.Series = s == s2 + check(assert_type(s > s2, "pd.Series[bool]"), pd.Series, np.bool_) + check(assert_type(s >= s2, "pd.Series[bool]"), pd.Series, np.bool_) + check(assert_type(s < s2, "pd.Series[bool]"), pd.Series, np.bool_) + check(assert_type(s <= s2, "pd.Series[bool]"), pd.Series, np.bool_) + check(assert_type(s == s2, "pd.Series[bool]"), pd.Series, np.bool_) def test_types_head_tail() -> None: @@ -309,7 +309,11 @@ def test_types_drop_multilevel() -> None: codes=[[0, 0, 0, 1, 1, 1], [0, 1, 2, 0, 1, 2]], ) s = pd.Series(data=[1, 2, 3, 4, 5, 6], index=index) - res: pd.Series = s.drop(labels="first", level=1) + check( + assert_type(s.drop(labels="first", level=1), "pd.Series[int]"), + pd.Series, + np.integer, + ) def test_types_drop_duplicates() -> None: @@ -382,7 +386,11 @@ def test_types_sort_index() -> None: # This was added in 1.1.0 https://pandas.pydata.org/docs/whatsnew/v1.1.0.html def test_types_sort_index_with_key() -> None: s = pd.Series([1, 2, 3], index=["a", "B", "c"]) - res: pd.Series = s.sort_index(key=lambda k: k.str.lower()) + check( + assert_type(s.sort_index(key=lambda k: k.str.lower()), "pd.Series[int]"), + pd.Series, + np.integer, + ) def test_types_sort_values() -> None: @@ -412,7 +420,11 @@ def test_types_sort_values() -> None: # This was added in 1.1.0 https://pandas.pydata.org/docs/whatsnew/v1.1.0.html def test_types_sort_values_with_key() -> None: s = pd.Series([1, 2, 3], index=[2, 3, 1]) - res: pd.Series = s.sort_values(key=lambda k: -k) + check( + assert_type(s.sort_values(key=lambda k: -k), "pd.Series[int]"), + pd.Series, + np.integer, + ) def test_types_shift() -> None: @@ -435,18 +447,32 @@ def test_types_rank() -> None: def test_types_mean() -> None: s = pd.Series([1, 2, 3, np.nan]) - f1: float = s.mean() - s1: pd.Series = s.groupby(level=0).mean() - f2: float = s.mean(skipna=False) - f3: float = s.mean(numeric_only=False) + check(assert_type(s.mean(), float), float) + check( + assert_type( + s.groupby(level=0).mean(), + "pd.Series[float]", + ), + pd.Series, + float, + ) + check(assert_type(s.mean(skipna=False), float), float) + check(assert_type(s.mean(numeric_only=False), float), float) def test_types_median() -> None: s = pd.Series([1, 2, 3, np.nan]) - f1: float = s.median() - s1: pd.Series = s.groupby(level=0).median() - f2: float = s.median(skipna=False) - f3: float = s.median(numeric_only=False) + check(assert_type(s.median(), float), float) + check( + assert_type( + s.groupby(level=0).median(), + "pd.Series[float]", + ), + pd.Series, + float, + ) + check(assert_type(s.median(skipna=False), float), float) + check(assert_type(s.median(numeric_only=False), float), float) def test_types_sum() -> None: @@ -620,67 +646,47 @@ def get_depth(url: str) -> int: check(assert_type(s.apply(lambda x: pd.NA), pd.Series), pd.Series, NAType) -def test_types_element_wise_arithmetic() -> None: - s = pd.Series([0, 1, -10]) - s2 = pd.Series([7, -5, 10]) - - res_add1: pd.Series = s + s2 - res_add2: pd.Series = s.add(s2, fill_value=0) - - res_sub: pd.Series = s - s2 - res_sub2: pd.Series = s.sub(s2, fill_value=0) - - res_mul: pd.Series = s * s2 - res_mul2: pd.Series = s.mul(s2, fill_value=0) - - res_div: pd.Series = s / s2 - res_div2: pd.Series = s.div(s2, fill_value=0) - - res_floordiv: pd.Series = s // s2 - res_floordiv2: pd.Series = s.floordiv(s2, fill_value=0) - - res_mod: pd.Series = s % s2 - res_mod2: pd.Series = s.mod(s2, fill_value=0) - - res_pow: pd.Series = s ** s2.abs() - res_pow2: pd.Series = s.pow(s2.abs(), fill_value=0) - - check(assert_type(divmod(s, s2), tuple["pd.Series[int]", "pd.Series[int]"]), tuple) - - def test_types_scalar_arithmetic() -> None: s = pd.Series([0, 1, -10]) - res_add1: pd.Series = s + 1 - res_add2: pd.Series = s.add(1, fill_value=0) + check(assert_type(s + 1, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.add(1, fill_value=0), "pd.Series[int]"), pd.Series, np.integer) - res_sub: pd.Series = s - 1 - res_sub2: pd.Series = s.sub(1, fill_value=0) + check(assert_type(s - 1, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.sub(1, fill_value=0), "pd.Series[int]"), pd.Series, np.integer) - res_mul: pd.Series = s * 2 - res_mul2: pd.Series = s.mul(2, fill_value=0) + check(assert_type(s * 2, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.mul(2, fill_value=0), "pd.Series[int]"), pd.Series, np.integer) - res_div: pd.Series = s / 2 - res_div2: pd.Series = s.div(2, fill_value=0) + check(assert_type(s / 2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.div(2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) - res_floordiv: pd.Series = s // 2 - res_floordiv2: pd.Series = s.floordiv(2, fill_value=0) + check(assert_type(s // 2, "pd.Series[int]"), pd.Series, np.integer) + check( + assert_type(s.floordiv(2, fill_value=0), "pd.Series[int]"), + pd.Series, + np.integer, + ) - res_mod: pd.Series = s % 2 - res_mod2: pd.Series = s.mod(2, fill_value=0) + check(assert_type(s % 2, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.mod(2, fill_value=0), "pd.Series[int]"), pd.Series, np.integer) - res_pow: pd.Series = s**2 - res_pow1: pd.Series = s**0 - res_pow2: pd.Series = s**0.213 - res_pow3: pd.Series = s.pow(0.5) + check(assert_type(s**2, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s**0, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s**0.213, "pd.Series[float]"), pd.Series, np.float64) + check(assert_type(s.pow(0.5), "pd.Series[float]"), pd.Series, np.float64) # GH 103 def test_types_complex_arithmetic() -> None: c = 1 + 1j s = pd.Series([1.0, 2.0, 3.0]) - x = s + c - y = s - c + check(assert_type(s + c, "pd.Series[complex]"), pd.Series, complex) + check(assert_type(s.add(c), "pd.Series[complex]"), pd.Series, complex) + check(assert_type(s - c, "pd.Series[complex]"), pd.Series, complex) + check(assert_type(s.sub(c), "pd.Series[complex]"), pd.Series, complex) def test_types_groupby() -> None: @@ -1105,8 +1111,9 @@ def test_types_getitem() -> None: s = pd.Series({"key": [0, 1, 2, 3]}) key: list[int] = s["key"] s2 = pd.Series([0, 1, 2, 3]) - value: int = s2[0] - s3: pd.Series = s[:2] + check(assert_type(s2[0], int), np.integer) + check(assert_type(s[:2], pd.Series), pd.Series) + check(assert_type(s2[:2], "pd.Series[int]"), pd.Series, np.integer) def test_types_getitem_by_timestamp() -> None: @@ -1117,9 +1124,9 @@ def test_types_getitem_by_timestamp() -> None: def test_types_eq() -> None: s1 = pd.Series([1, 2, 3]) - res1: pd.Series = s1 == 1 + check(assert_type(s1 == 1, "pd.Series[bool]"), pd.Series, np.bool_) s2 = pd.Series([1, 2, 4]) - res2: pd.Series = s1 == s2 + check(assert_type(s1 == s2, "pd.Series[bool]"), pd.Series, np.bool_) def test_types_rename_axis() -> None: @@ -1263,10 +1270,12 @@ def test_types_dot() -> None: s1 = pd.Series([0, 1, 2, 3]) s2 = pd.Series([-1, 2, -3, 4]) df1 = pd.DataFrame([[0, 1], [-2, 3], [4, -5], [6, 7]]) + df2 = pd.DataFrame([[0.0, 1.0], [-2.0, 3.0], [4.0, -5.0], [6.0, 7.0]]) n1 = np.array([[0, 1], [1, 2], [-1, -1], [2, 0]]) - check(assert_type(s1.dot(s2), Scalar), np.int64) - check(assert_type(s1 @ s2, Scalar), np.int64) - check(assert_type(s1.dot(df1), "pd.Series[int]"), pd.Series, np.int64) + check(assert_type(s1.dot(s2), Scalar), np.integer) + check(assert_type(s1 @ s2, Scalar), np.integer) + check(assert_type(s1.dot(df1), pd.Series), pd.Series, np.integer) + check(assert_type(s1.dot(df2), pd.Series), pd.Series, np.float64) check(assert_type(s1 @ df1, pd.Series), pd.Series) check(assert_type(s1.dot(n1), np.ndarray), np.ndarray) check(assert_type(s1 @ n1, np.ndarray), np.ndarray) @@ -1289,7 +1298,7 @@ def test_series_min_max_sub_axis() -> None: check(assert_type(sa, pd.Series), pd.Series) check(assert_type(ss, pd.Series), pd.Series) check(assert_type(sm, pd.Series), pd.Series) - check(assert_type(sd, pd.Series), pd.Series) + check(assert_type(sd, "pd.Series[float]"), pd.Series) def test_series_index_isin() -> None: @@ -1317,17 +1326,17 @@ def test_series_multiindex_getitem() -> None: s = pd.Series( [1, 2, 3, 4], index=pd.MultiIndex.from_product([["a", "b"], ["x", "y"]]) ) - s1: pd.Series = s["a", :] + check(assert_type(s["a", :], "pd.Series[int]"), pd.Series, np.integer) def test_series_mul() -> None: s = pd.Series([1, 2, 3]) sm = s * 4 - check(assert_type(sm, pd.Series), pd.Series) + check(assert_type(sm, "pd.Series[int]"), pd.Series, np.integer) ss = s - 4 - check(assert_type(ss, pd.Series), pd.Series) + check(assert_type(ss, "pd.Series[int]"), pd.Series, np.integer) sm2 = s * s - check(assert_type(sm2, pd.Series), pd.Series) + check(assert_type(sm2, "pd.Series[int]"), pd.Series, np.integer) sp = s + 4 check(assert_type(sp, "pd.Series[int]"), pd.Series, np.integer) @@ -1408,13 +1417,19 @@ def test_cat_accessor() -> None: def test_cat_ctor_values() -> None: - c1 = pd.Categorical(["a", "b", "a"]) + check(assert_type(pd.Categorical(["a", "b", "a"]), pd.Categorical), pd.Categorical) # GH 95 - c2 = pd.Categorical(pd.Series(["a", "b", "a"])) + check( + assert_type(pd.Categorical(pd.Series(["a", "b", "a"])), pd.Categorical), + pd.Categorical, + ) s: Sequence = cast(Sequence, ["a", "b", "a"]) - c3 = pd.Categorical(s) + check(assert_type(pd.Categorical(s), pd.Categorical), pd.Categorical) # GH 107 - c4 = pd.Categorical(np.array([1, 2, 3, 1, 1])) + check( + assert_type(pd.Categorical(np.array([1, 2, 3, 1, 1])), pd.Categorical), + pd.Categorical, + ) def test_iloc_getitem_ndarray() -> None: @@ -2768,7 +2783,7 @@ def test_astype_other() -> None: def test_all_astype_args_tested() -> None: """Check that all relevant numpy type aliases are tested.""" - NUMPY_ALIASES: set[str] = {k for k in np.sctypeDict} + NUMPY_ALIASES: set[str | int] = {k for k in np.sctypeDict} EXCLUDED_ALIASES = { "datetime64", "m", @@ -2889,7 +2904,7 @@ def test_convert_dtypes_dtype_backend() -> None: def test_apply_returns_none() -> None: # GH 557 s = pd.Series([1, 2, 3]) - check(assert_type(s.apply(lambda x: None), pd.Series), pd.Series) + check(assert_type(s.apply(lambda _: None), pd.Series), pd.Series) def test_loc_callable() -> None: diff --git a/tests/test_series_arithmetic.py b/tests/test_series_arithmetic.py new file mode 100644 index 000000000..59b728da3 --- /dev/null +++ b/tests/test_series_arithmetic.py @@ -0,0 +1,281 @@ +"""Test module for arithmetic operations on Series.""" + +import numpy as np +import pandas as pd +from typing_extensions import assert_type + +from tests import check + + +def test_element_wise_int_int() -> None: + s = pd.Series([0, 1, -10]) + s2 = pd.Series([7, -5, 10]) + + check(assert_type(s + s2, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.add(s2, fill_value=0), "pd.Series[int]"), pd.Series, np.integer) + + check(assert_type(s - s2, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.sub(s2, fill_value=0), "pd.Series[int]"), pd.Series, np.integer) + + check(assert_type(s * s2, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.mul(s2, fill_value=0), "pd.Series[int]"), pd.Series, np.integer) + + # GH1089 should be the following + check(assert_type(s / s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.div(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s // s2, "pd.Series[int]"), pd.Series, np.integer) + check( + assert_type(s.floordiv(s2, fill_value=0), "pd.Series[int]"), + pd.Series, + np.integer, + ) + + check(assert_type(s % s2, "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.mod(s2, fill_value=0), "pd.Series[int]"), pd.Series, np.integer) + + check(assert_type(s ** s2.abs(), "pd.Series[int]"), pd.Series, np.integer) + check( + assert_type(s.pow(s2.abs(), fill_value=0), "pd.Series[int]"), + pd.Series, + np.integer, + ) + + check(assert_type(divmod(s, s2), tuple["pd.Series[int]", "pd.Series[int]"]), tuple) + + +def test_element_wise_int_float() -> None: + s = pd.Series([0, 1, -10]) + s2 = pd.Series([7.0, -5.5, 10.4]) + + check(assert_type(s + s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.add(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s - s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.sub(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s * s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.mul(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + # GH1089 should be the following + check(assert_type(s / s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.div(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check( + assert_type(s.floordiv(s2, fill_value=0), "pd.Series[float]"), + pd.Series, + np.float64, + ) + + check(assert_type(s % s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.mod(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s ** s2.abs(), "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.pow(s2.abs(), fill_value=0), "pd.Series[float]"), + pd.Series, + np.float64, + ) + + check(assert_type(divmod(s, s2), tuple["pd.Series[int]", "pd.Series[int]"]), tuple) + + +def test_element_wise_float_int() -> None: + s = pd.Series([0.0, 1.4, -10.25]) + s2 = pd.Series([7, -5, 10]) + + check(assert_type(s + s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.add(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s - s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.sub(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s * s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.mul(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + # GH1089 should be the following + check(assert_type(s / s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.div(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s // s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.floordiv(s2, fill_value=0), "pd.Series[float]"), + pd.Series, + np.float64, + ) + + check(assert_type(s % s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.mod(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s ** s2.abs(), "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.pow(s2.abs(), fill_value=0), "pd.Series[float]"), + pd.Series, + np.float64, + ) + + check( + assert_type(divmod(s, s2), tuple["pd.Series[float]", "pd.Series[float]"]), tuple + ) + + +def test_element_wise_float_float() -> None: + s = pd.Series([0.439, 1.43829, -10.432]) + s2 = pd.Series([7.4389, -5.543, 10.432]) + + check(assert_type(s + s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.add(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s - s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.sub(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s * s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.mul(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + # GH1089 should be the following + check(assert_type(s / s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.div(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s // s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.floordiv(s2, fill_value=0), "pd.Series[float]"), + pd.Series, + np.float64, + ) + + check(assert_type(s % s2, "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.mod(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64 + ) + + check(assert_type(s ** s2.abs(), "pd.Series[float]"), pd.Series, np.float64) + check( + assert_type(s.pow(s2.abs(), fill_value=0), "pd.Series[float]"), + pd.Series, + np.float64, + ) + + check( + assert_type(divmod(s, s2), tuple["pd.Series[float]", "pd.Series[float]"]), tuple + ) + + +def test_element_wise_int_unknown() -> None: + 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), 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) + + # 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)