如何将pandas.DataFrame转换为二维数组？（pytest场景）

问题描述 #

通过代码tests = pd.read_csv('tsur.tsv', sep='\t')读取得到如下DataFrame：

query  match
0                                        jordan belfort  False
1       Ricardo Jalone Bob FRANK\n\n1\nLondon Markus  False
2                                        USA Maria JOHN  False

需要将其转换为如下格式的元组列表：

[('jordan belfort', False), ('Ricardo Jalone Bob FRANK\n\n1\nLondon Markus', False), ('USA Maria JOHN', False)]

并将该列表集成到pytest的fixture中。

实现步骤 #

1. DataFrame转换为目标格式 #

利用pandas的itertuples()方法可直接将每行数据转换为元组，再转为列表：

import pandas as pd

def load_test_data():
    df = pd.read_csv('tsur.tsv', sep='\t')
    # index=False 排除默认行索引，只保留query和match列内容
    return list(df.itertuples(index=False, name=None))

调用该函数即可得到所需的元组列表。

2. 定义pytest fixture #

在测试文件（或conftest.py中，方便全局复用）中定义fixture：

import pytest
import pandas as pd

@pytest.fixture
def test_data():
    df = pd.read_csv('tsur.tsv', sep='\t')
    return list(df.itertuples(index=False, name=None))

3. 在测试用例中使用fixture #

直接在测试函数参数中传入fixture名称即可调用：

def test_something(test_data):
    # 遍历测试数据执行逻辑
    for query, expected_match in test_data:
        result = your_test_function(query)
        assert result == expected_match

内容的提问来源于stack exchange，提问作者reader074