12. Flatten a List / Вирішуємо задачі на Python CheckIO Українською

# ----------   tasks/flatten.py   -----------


def is_leaf(array):
    return isinstance(array, int)


def flat_list(array):
    res = []

    for i in array:
        if is_leaf(i):
            res.append(i)
        else:
            leafs = flat_list(i)
            res += leafs

    return res


def flat_list_while(array):
    res = []

    query = list(array)

    while query:
        i = query.pop(0)
        if is_leaf(i):
            res.append(i)
        else:
            items = list(i)
            query = items + query

    return res



# ------------  tests/test_flatten.py  -----------------

import pytest

from tasks.flatten import flat_list, flat_list_while


@pytest.mark.parametrize("func", [flat_list, flat_list_while])
@pytest.mark.parametrize(["data", "res"], [
    [[], []],
    [[1, 2, 3], [1, 2, 3]],
    [[1, [2, 2, 2], 4], [1, 2, 2, 2, 4]],
    [[[[2]], [4, [5, 6, [6], 6, 6, 6], 7]], [2, 4, 5, 6, 6, 6, 6, 6, 7]],
    [[-1, [1, [-2], 1], -1], [-1, 1, -2, 1, -1]]
])
def test_flatten(func, data, res):
    assert func(data) == res