Python itertools example

combine_lists.py

# run tests uisng: pytest -vv combine_lists.py 

from typing import List, Any
import itertools


def combine_lists_of_lists(*args: List[List[Any]]) -> List[List[Any]]:
    results = [list(itertools.chain(*tup))
               for tup in itertools.zip_longest(*args, fillvalue=[])]
    return results


def test_1():
    list1 = [["a"], [1]]
    list2 = [["b"], [2]]
    assert combine_lists_of_lists(list1, list2) == [["a", "b"], [1, 2]]


def test_2():
    list1 = [["a"], [1]]
    list2 = []
    list3 = [["b"], [2]]
    assert combine_lists_of_lists(list1, list2, list3) == [["a", "b"], [1, 2]]


def test_3():
    list1 = [["a", "b"], [1, 2]]
    list2 = []
    list3 = [["c"], [3, 4]]
    assert combine_lists_of_lists(list1, list2, list3) == [["a", "b", "c"], [1, 2, 3, 4]]



def combine_lists_of_lists_original(*args):
    results = []
    for rows in args:
        for i, x in enumerate(rows):
            if len(results) <= i:
                results.append(x)
            else:
                results[i].extend(x)


    return results