test_utils.py 文件源码

def test_remove_missing():
    df = pd.DataFrame({'a': [1.0, np.NaN, 3, np.inf],
                       'b': [1, 2, 3, 4]})
    df2 = pd.DataFrame({'a': [1.0, 3, np.inf],
                       'b': [1, 3, 4]})
    df3 = pd.DataFrame({'a': [1.0, 3],
                       'b': [1, 3]})

    with warnings.catch_warnings(record=True) as w:
        res = remove_missing(df, na_rm=True, vars=['b'])
        res.equals(df)

        res = remove_missing(df)
        res.equals(df2)

        res = remove_missing(df, na_rm=True, finite=True)
        res.equals(df3)
        assert len(w) == 1