Expands a list of lists as a tree diagram

expand_tree_diagram.py

def expand_tree_diagram(a_list, a_branch = []):
    if a_list == []:
        return [a_branch]
    else:
        result = []
        for a_leaf in a_list[0]:
            result += expand_tree_diagram(a_list[1:], a_branch + [a_leaf])
    return result

if __name__ == '__main__':
    t = expand_tree_diagram([[1,2,3,4],["a", "b"],["i", "ii", "iii", "iv"], ["A", "B", "C"]])
    assert t == [[1, 'a', 'i', 'A'], [1, 'a', 'i', 'B'], [1, 'a', 'i', 'C'], [1, 'a', 'ii', 'A'], [1, 'a', 'ii', 'B'], [1, 'a', 'ii', 'C'], [1, 'a', 'iii', 'A'], [1, 'a', 'iii', 'B'], [1, 'a', 'iii', 'C'], [1, 'a', 'iv', 'A'], [1, 'a', 'iv', 'B'], [1, 'a', 'iv', 'C'], [1, 'b', 'i', 'A'], [1, 'b', 'i', 'B'], [1, 'b', 'i', 'C'], [1, 'b', 'ii', 'A'], [1, 'b', 'ii', 'B'], [1, 'b', 'ii', 'C'], [1, 'b', 'iii', 'A'], [1, 'b', 'iii', 'B'], [1, 'b', 'iii', 'C'], [1, 'b', 'iv', 'A'], [1, 'b', 'iv', 'B'], [1, 'b', 'iv', 'C'], [2, 'a', 'i', 'A'], [2, 'a', 'i', 'B'], [2, 'a', 'i', 'C'], [2, 'a', 'ii', 'A'], [2, 'a', 'ii', 'B'], [2, 'a', 'ii', 'C'], [2, 'a', 'iii', 'A'], [2, 'a', 'iii', 'B'], [2, 'a', 'iii', 'C'], [2, 'a', 'iv', 'A'], [2, 'a', 'iv', 'B'], [2, 'a', 'iv', 'C'], [2, 'b', 'i', 'A'], [2, 'b', 'i', 'B'], [2, 'b', 'i', 'C'], [2, 'b', 'ii', 'A'], [2, 'b', 'ii', 'B'], [2, 'b', 'ii', 'C'], [2, 'b', 'iii', 'A'], [2, 'b', 'iii', 'B'], [2, 'b', 'iii', 'C'], [2, 'b', 'iv', 'A'], [2, 'b', 'iv', 'B'], [2, 'b', 'iv', 'C'], [3, 'a', 'i', 'A'], [3, 'a', 'i', 'B'], [3, 'a', 'i', 'C'], [3, 'a', 'ii', 'A'], [3, 'a', 'ii', 'B'], [3, 'a', 'ii', 'C'], [3, 'a', 'iii', 'A'], [3, 'a', 'iii', 'B'], [3, 'a', 'iii', 'C'], [3, 'a', 'iv', 'A'], [3, 'a', 'iv', 'B'], [3, 'a', 'iv', 'C'], [3, 'b', 'i', 'A'], [3, 'b', 'i', 'B'], [3, 'b', 'i', 'C'], [3, 'b', 'ii', 'A'], [3, 'b', 'ii', 'B'], [3, 'b', 'ii', 'C'], [3, 'b', 'iii', 'A'], [3, 'b', 'iii', 'B'], [3, 'b', 'iii', 'C'], [3, 'b', 'iv', 'A'], [3, 'b', 'iv', 'B'], [3, 'b', 'iv', 'C'], [4, 'a', 'i', 'A'], [4, 'a', 'i', 'B'], [4, 'a', 'i', 'C'], [4, 'a', 'ii', 'A'], [4, 'a', 'ii', 'B'], [4, 'a', 'ii', 'C'], [4, 'a', 'iii', 'A'], [4, 'a', 'iii', 'B'], [4, 'a', 'iii', 'C'], [4, 'a', 'iv', 'A'], [4, 'a', 'iv', 'B'], [4, 'a', 'iv', 'C'], [4, 'b', 'i', 'A'], [4, 'b', 'i', 'B'], [4, 'b', 'i', 'C'], [4, 'b', 'ii', 'A'], [4, 'b', 'ii', 'B'], [4, 'b', 'ii', 'C'], [4, 'b', 'iii', 'A'], [4, 'b', 'iii', 'B'], [4, 'b', 'iii', 'C'], [4, 'b', 'iv', 'A'], [4, 'b', 'iv', 'B'], [4, 'b', 'iv', 'C']]