diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 457db8de3..eb980e79f 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -798,8 +798,18 @@ class Series(IndexOpsMixin[S1], NDFrame): def dot( self, other: ArrayLike | dict[_str, np.ndarray] | Sequence[S1] | Index[S1] ) -> np.ndarray: ... - def __matmul__(self, other): ... - def __rmatmul__(self, other): ... + @overload + def __matmul__(self, other: Series) -> Scalar: ... + @overload + def __matmul__(self, other: DataFrame) -> Series: ... + @overload + def __matmul__(self, other: np.ndarray) -> np.ndarray: ... + @overload + def __rmatmul__(self, other: Series) -> Scalar: ... + @overload + def __rmatmul__(self, other: DataFrame) -> Series: ... + @overload + def __rmatmul__(self, other: np.ndarray) -> np.ndarray: ... @overload def searchsorted( self, diff --git a/tests/test_series.py b/tests/test_series.py index 8e9b60bde..532be0625 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -1236,16 +1236,17 @@ def test_types_as_type() -> None: def test_types_dot() -> None: + """Test typing of multiplication methods (dot and @) for Series.""" 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]]) n1 = np.array([[0, 1], [1, 2], [-1, -1], [2, 0]]) - sc1: Scalar = s1.dot(s2) - sc2: Scalar = s1 @ s2 - s3: pd.Series = s1.dot(df1) - s4: pd.Series = s1 @ df1 - n2: np.ndarray = s1.dot(n1) - n3: np.ndarray = s1 @ n1 + 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 @ df1, pd.Series), pd.Series) + check(assert_type(s1.dot(n1), np.ndarray), np.ndarray) + check(assert_type(s1 @ n1, np.ndarray), np.ndarray) def test_series_loc_setitem() -> None: