From e95d425e1afc62b8a7d344d8c1d2d53492bf46c2 Mon Sep 17 00:00:00 2001 From: Amanda Bizzinotto Date: Tue, 6 Jun 2023 19:14:18 -0300 Subject: [PATCH 1/3] Update series.pyi --- pandas-stubs/core/series.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 = ..., From b89d5fa67c6e6f02749bab7705e042a2b3e192b7 Mon Sep 17 00:00:00 2001 From: Amanda Bizzinotto Date: Tue, 6 Jun 2023 19:34:50 -0300 Subject: [PATCH 2/3] Add type tests for mark --- tests/test_series.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/test_series.py b/tests/test_series.py index b34fd039f..3a7dd529c 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.int64) + + # 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.int64) + + # 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) From 67dbaa7da70f64c3fe2938f9dab45dc8531abb66 Mon Sep 17 00:00:00 2001 From: Amanda Bizzinotto Date: Wed, 7 Jun 2023 12:16:36 -0300 Subject: [PATCH 3/3] Use np.integer instead of np.int64 --- tests/test_series.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_series.py b/tests/test_series.py index 3a7dd529c..c8cc0d65f 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -2020,13 +2020,13 @@ 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.int64) + 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.int64) + 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)