Pandas DataFrame中NA值的单元测试问题求助

背景信息 #

有一个生成Pandas DataFrame的函数，输出的DataFrame结构如下：

各列数据类型：

timestamp      Int64
duration       Int64
trial_type    object
blink          Int64
message       object
dtype: object

编写的单元测试代码如下，用于验证第一行输出：

@pytest.mark.parametrize(     
    "folder, expected",     
    [("emg", [9199380, pd.NA, np.nan, pd.NA, "RECORD_START"])]
)
def test_physioevents_value(folder, expected, eyelink_test_data_dir):
    input_dir = eyelink_test_data_dir / folder
    asc_file = asc_test_files(input_dir=input_dir, suffix="*_events")[0]
    events = _load_asc_file(asc_file)
    events_after_start = _df_events_after_start(events)
    physioevents_reordered = _df_physioevents(events_after_start)
    physioevents_eye1 = _physioevents_eye1(physioevents_reordered)
    assert physioevents_eye1.iloc[0].tolist() == expected

测试报错情况 #

运行测试时出现如下错误：

E AssertionError: assert [9199380, ...CORD_START'] == [9199380, ...CORD_START']
E
E (pytest_assertion plugin: representation of details failed: missing.pyx:392: TypeError: boolean value of NA is ambiguous.
E Probably an object has a faulty repr.)

tests/test_edf2bids.py:670: AssertionError

直接用pd.NA或字符串""都无法匹配DataFrame中的空值，且不希望将所有列转为object类型（会破坏数值列的类型，比如blink列会变成float）。

解决方案 #

方法1：使用pd.testing.assert_series_equal进行精确比较 #

Pandas提供了专门用于Series比较的工具方法，能正确处理不同类型的空值（pd.NA、np.nan）以及数据类型匹配问题。替换原来的assert语句：

import pandas as pd

# 替换原assert行
expected_series = pd.Series(expected, index=physioevents_eye1.columns)
pd.testing.assert_series_equal(physioevents_eye1.iloc[0], expected_series, check_dtype=True)

如果不需要严格匹配数据类型，可以添加check_dtype=False参数；如果不需要匹配索引名称，添加check_names=False。

方法2：手动遍历元素进行空值安全比较 #

自定义比较逻辑，用pd.isna()判断空值，避免直接用==比较空值（因为pd.NA == pd.NA会返回pd.NA，导致布尔值歧义）：

actual = physioevents_eye1.iloc[0].tolist()
for a, e in zip(actual, expected):
    if pd.isna(a) and pd.isna(e):
        continue
    assert a == e, f"Actual value {a} != Expected value {e}"

方法3：统一expected中的空值类型 #

确保expected中的空值类型与DataFrame列类型完全匹配：

• 对于Int64类型列（timestamp、duration、blink）：用pd.NA
• 对于object类型字符串列（trial_type、message）：用np.nan

调整后的parametrize部分：

@pytest.mark.parametrize(     
    "folder, expected",     
    [("emg", [9199380, pd.NA, np.nan, pd.NA, "RECORD_START"])]
)

注：此方法需配合方法1或方法2使用，直接列表比较仍会报错，因为列表中的pd.NA和DataFrame中的pd.NA直接比较仍有歧义。

内容的提问来源于stack exchange，提问作者Julia-Katharina Pfarr