|
1 | 1 | import numpy as np
|
2 | 2 | import pytest
|
3 | 3 |
|
4 |
| -from pandas import Series |
| 4 | +from pandas import ( |
| 5 | + Int64Dtype, |
| 6 | + Series, |
| 7 | +) |
5 | 8 | import pandas._testing as tm
|
6 | 9 |
|
7 | 10 |
|
@@ -71,15 +74,8 @@ def test_mask_inplace():
|
71 | 74 |
|
72 | 75 | def test_mask_na():
|
73 | 76 | # We should not be filling pd.NA. See GH#60729
|
74 |
| - series = Series([None, 1, 2, None, 3, 4, None]) |
75 |
| - series = series.convert_dtypes() |
76 |
| - cond = series <= 2 |
| 77 | + series = Series([None, 1, 2, None, 3, 4, None], dtype=Int64Dtype()) |
| 78 | + result = series.mask(series <= 2, -99) |
| 79 | + expected = Series([None, 1, 2, None, -99, -99, None], dtype=Int64Dtype()) |
77 | 80 |
|
78 |
| - maskres = series.mask(cond, -99) |
79 |
| - whereres = series.where(~(cond), -99) |
80 |
| - |
81 |
| - expected = Series([None, -99, -99, None, 3, 4, None]) |
82 |
| - expected = expected.convert_dtypes() |
83 |
| - |
84 |
| - tm.assert_series_equal(maskres, expected) |
85 |
| - tm.assert_series_equal(maskres, whereres) |
| 81 | + tm.assert_series_equal(result, expected) |
0 commit comments