Search code examples
pythonpython-3.xfor-loopnested-loopspython-itertools

How can I use itertools to simplify this nested for loop?


How can I use itertools to simplify this nested for loop?

# 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

Sample Data is in 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"

Explanation

The YAML file contains parameters uses for building test cases. We have a block of test cases for Cisco, and other test cases for Juniper. So the structure is this vendor --> name of class --> function --> test case data

Function get_test_cases in my pytest file will load the YAML file contents, and then iterates through the contents, and yield data for each test case. Function test_network consumes this data, and performs the actual testing.

I am trying to find an alternative to the nested for loops so that I can raise my pylint rating... if this is even possible.


Solution

  • You could try the following to make it a bit compacter:

    def testcases(data, level=0):
        if level == 3:
            for case in data:
                yield (case,)
        else:
            for record in data:
                for name in record:
                    for case in testcases(record[name], level + 1):
                        yield (name,) + case
    

    or

    def testcases(data, level=0):
        if level == 6:
            for case in data:
                yield (case,)
        elif level % 2 == 0:
            for record in data:
                yield from testcases(record, level + 1)
        else:
            for name, records in data.items():
                for case in testcases(records, level + 1):
                    yield (name,) + case
    

    and then

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

    With your sample file this

    for case in get_test_cases():
        print(case)
    

    does produce

    ('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'})