86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
"""
|
|
Tests for DataFrame.mask; tests DataFrame.where as a side-effect.
|
|
"""
|
|
|
|
import numpy as np
|
|
|
|
from pandas import DataFrame, isna
|
|
import pandas._testing as tm
|
|
|
|
|
|
class TestDataFrameMask:
|
|
def test_mask(self):
|
|
df = DataFrame(np.random.randn(5, 3))
|
|
cond = df > 0
|
|
|
|
rs = df.where(cond, np.nan)
|
|
tm.assert_frame_equal(rs, df.mask(df <= 0))
|
|
tm.assert_frame_equal(rs, df.mask(~cond))
|
|
|
|
other = DataFrame(np.random.randn(5, 3))
|
|
rs = df.where(cond, other)
|
|
tm.assert_frame_equal(rs, df.mask(df <= 0, other))
|
|
tm.assert_frame_equal(rs, df.mask(~cond, other))
|
|
|
|
# see GH#21891
|
|
df = DataFrame([1, 2])
|
|
res = df.mask([[True], [False]])
|
|
|
|
exp = DataFrame([np.nan, 2])
|
|
tm.assert_frame_equal(res, exp)
|
|
|
|
def test_mask_inplace(self):
|
|
# GH#8801
|
|
df = DataFrame(np.random.randn(5, 3))
|
|
cond = df > 0
|
|
|
|
rdf = df.copy()
|
|
|
|
return_value = rdf.where(cond, inplace=True)
|
|
assert return_value is None
|
|
tm.assert_frame_equal(rdf, df.where(cond))
|
|
tm.assert_frame_equal(rdf, df.mask(~cond))
|
|
|
|
rdf = df.copy()
|
|
return_value = rdf.where(cond, -df, inplace=True)
|
|
assert return_value is None
|
|
tm.assert_frame_equal(rdf, df.where(cond, -df))
|
|
tm.assert_frame_equal(rdf, df.mask(~cond, -df))
|
|
|
|
def test_mask_edge_case_1xN_frame(self):
|
|
# GH#4071
|
|
df = DataFrame([[1, 2]])
|
|
res = df.mask(DataFrame([[True, False]]))
|
|
expec = DataFrame([[np.nan, 2]])
|
|
tm.assert_frame_equal(res, expec)
|
|
|
|
def test_mask_callable(self):
|
|
# GH#12533
|
|
df = DataFrame([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
|
result = df.mask(lambda x: x > 4, lambda x: x + 1)
|
|
exp = DataFrame([[1, 2, 3], [4, 6, 7], [8, 9, 10]])
|
|
tm.assert_frame_equal(result, exp)
|
|
tm.assert_frame_equal(result, df.mask(df > 4, df + 1))
|
|
|
|
# return ndarray and scalar
|
|
result = df.mask(lambda x: (x % 2 == 0).values, lambda x: 99)
|
|
exp = DataFrame([[1, 99, 3], [99, 5, 99], [7, 99, 9]])
|
|
tm.assert_frame_equal(result, exp)
|
|
tm.assert_frame_equal(result, df.mask(df % 2 == 0, 99))
|
|
|
|
# chain
|
|
result = (df + 2).mask(lambda x: x > 8, lambda x: x + 10)
|
|
exp = DataFrame([[3, 4, 5], [6, 7, 8], [19, 20, 21]])
|
|
tm.assert_frame_equal(result, exp)
|
|
tm.assert_frame_equal(result, (df + 2).mask((df + 2) > 8, (df + 2) + 10))
|
|
|
|
def test_mask_dtype_conversion(self):
|
|
# GH#3733
|
|
df = DataFrame(data=np.random.randn(100, 50))
|
|
df = df.where(df > 0) # create nans
|
|
bools = df > 0
|
|
mask = isna(df)
|
|
expected = bools.astype(float).mask(mask)
|
|
result = bools.mask(mask)
|
|
tm.assert_frame_equal(result, expected)
|