Skip to content

Commit d2d5f62

Browse files
sanggon6107WillAyd
andcommitted
ENH: Optimze test_mask_na()
Co-authored-by: WillAyd <[email protected]>
1 parent 6fd8986 commit d2d5f62

File tree

1 file changed

+8
-12
lines changed

1 file changed

+8
-12
lines changed

pandas/tests/series/indexing/test_mask.py

+8-12
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
import numpy as np
22
import pytest
33

4-
from pandas import Series
4+
from pandas import (
5+
Int64Dtype,
6+
Series,
7+
)
58
import pandas._testing as tm
69

710

@@ -71,15 +74,8 @@ def test_mask_inplace():
7174

7275
def test_mask_na():
7376
# 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())
7780

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

Comments
 (0)