如何使用itertools简化pytest测试用例生成中的嵌套for循环？

问题背景 #

你编写的get_test_cases函数通过多层嵌套for循环遍历YAML测试数据生成测试用例参数，但多层嵌套导致pylint评分较低，希望借助itertools工具简化代码结构。

原代码与YAML结构回顾 #

原pytest代码 #

# pytest file:

def get_test_cases():

    with open("file_path.yml", "r", encoding="utf-8") as index_file:
        data = yaml.safe_load(index_file)

    for vendor_dict in data:
        for vendor_name, class_list in vendor_dict.items():
            for class_dict in class_list:
                for class_name, method_list in class_dict.items():
                    for method_dict in method_list:
                        for function, test_list in method_dict.items():
                            for test_case in test_list:
                                yield vendor_name, class_name, function, test_case

@pytest.mark.parametrize("iteration", get_test_cases())
def test_network(iteration, monkeypatch):
    """Performs pytest using info provided from test case"""

    (vendor_name, class_name, function, test_case) = iteration

    # Use data above to perform pytest
    # Code omitted

示例YAML数据 #

# YAML file:

- cisco:
    - CiscoClass:
        - get_interface_stats:
            - test_description: "test interface stats"
              function_input: "Gig1/1/1"
              expected_output: "show interfaces Gig1/1/1"

- juniper:
    - JuniperClass:
        - get_vlan_info: 
            - test_description: "test get_vlan_info"
              function_input: "10"
              expected_output: "show vlan 10"

YAML结构层级为：厂商 → 类列表 → 方法列表 → 测试用例列表，每一层都是单键字典的列表形式。

解决方案：用itertools扁平化嵌套层级 #

利用itertools.chain.from_iterable逐层拆解嵌套结构，替代多层for循环，降低代码圈复杂度，提升pylint评分。

优化后的代码 #

import itertools
import yaml
import pytest

def get_test_cases():
    with open("file_path.yml", "r", encoding="utf-8") as index_file:
        data = yaml.safe_load(index_file)

    # 1. 扁平化厂商层：提取所有(vendor_name, class_list)键值对
    vendor_pairs = itertools.chain.from_iterable(vd.items() for vd in data)
    
    # 2. 扁平化类层：提取所有(vendor_name, class_name, method_list)元组
    class_pairs = itertools.chain.from_iterable(
        ((vendor_name, *cd.items()) for cd in class_list)
        for vendor_name, class_list in vendor_pairs
    )
    
    # 3. 扁平化方法层：提取所有(vendor_name, class_name, function, test_list)元组
    method_pairs = itertools.chain.from_iterable(
        ((vendor_name, class_name, *md.items()) for md in method_list)
        for vendor_name, class_name, method_list in class_pairs
    )
    
    # 4. 展开测试用例层：生成最终的测试用例参数
    for vendor_name, class_name, function, test_list in method_pairs:
        yield from ((vendor_name, class_name, function, tc) for tc in test_list)

@pytest.mark.parametrize("vendor_name, class_name, function, test_case", get_test_cases())
def test_network(vendor_name, class_name, function, test_case, monkeypatch):
    """Performs pytest using info provided from test case"""
    # 直接使用解包后的参数，无需再手动拆分迭代器
    # Code omitted

代码说明 #

1. 厂商层扁平化：用itertools.chain.from_iterable把每个厂商字典的键值对（单键）合并成一个迭代器，避免第一层循环。
2. 类层扁平化：遍历每个厂商对应的类列表，用*cd.items()解包单键字典的键值对，生成包含厂商名、类名、方法列表的元组，再用chain.from_iterable合并。
3. 方法层扁平化：重复类层的逻辑，生成包含厂商名、类名、函数名、测试用例列表的元组。
4. 测试用例展开：用yield from直接生成每个测试用例的最终参数组，替代最后一层循环。
5. 测试函数优化：把参数直接解包到测试函数的形参中，省去手动拆分迭代器的步骤，代码更简洁。

另一种简洁写法（生成器表达式） #

如果不需要严格依赖itertools，也可以把多层嵌套循环转换成单一生成器表达式，同样能降低圈复杂度：

def get_test_cases():
    with open("file_path.yml", "r", encoding="utf-8") as index_file:
        data = yaml.safe_load(index_file)

    return (
        (vendor_name, class_name, function, test_case)
        for vendor_dict in data
        for vendor_name, class_list in vendor_dict.items()
        for class_dict in class_list
        for class_name, method_list in class_dict.items()
        for method_dict in method_list
        for function, test_list in method_dict.items()
        for test_case in test_list
    )

这种写法把嵌套循环合并成一行生成器表达式，pylint会认为其圈复杂度低于多层独立for循环，同时保持逻辑清晰。

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