diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index ae6daa5da..ec98f9b9f 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1378,7 +1378,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def mask( self, cond: MaskType, - other: Scalar | Series[S1] | DataFrame | Callable = ..., + other: Scalar | Series[S1] | DataFrame | Callable | NAType | None = ..., *, inplace: _bool = ..., axis: AxisIndex | None = ..., diff --git a/tests/test_series.py b/tests/test_series.py index b34fd039f..c8cc0d65f 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -2014,3 +2014,20 @@ def test_to_string() -> None: ), str, ) + + +def test_types_mask() -> None: + s = pd.Series([1, 2, 3, 4, 5]) + + # Test case with a boolean condition and a scalar value + check(assert_type(s.mask(s > 3, 10), pd.Series), pd.Series, np.integer) + + # Test case with a boolean condition and a callable + def double(x): + return x * 2 + + check(assert_type(s.mask(s > 3, double), pd.Series), pd.Series, np.integer) + + # Test cases with None and pd.NA as other + check(assert_type(s.mask(s > 3, None), pd.Series), pd.Series, np.float64) + check(assert_type(s.mask(s > 3, pd.NA), pd.Series), pd.Series, np.float64)