Partial list Flattening in Python

Question

I have a small problem where I lack quite some Python experience. Assuming a list:

list=[[40, 20, 45, 40], [[20, 10, 10, 30, 45, 20], [30, 20, 20, 30]]]

I want to do something I call 'partial flattening' because the expected output is:

[[40, 20, 45, 40], [20, 10, 10, 30, 45, 20], [30, 20, 20, 30]]

Notice before the lentgh of list was 2 while the expected output has length of 3 . The problem is that I don't know in advance which element from list is nested one layer deeper.

This answer and many others didn't help much, because the flattening goes to far. Following the question the output is:

list=[40, 20, 45, 40, [20, 10, 10, 30, 45, 20], [30, 20, 20, 30]]

(notice the missing brackets for the first 4 elements from the list.

Answer 1

You can create a recursive function, testing whether all the inner elements are lists. If so, recursively apply the algorithm, otherwise yield the list itself.

def inner_lists(lst):
    if all(isinstance(x, list) for x in lst):
        for inner in lst:
            for x in inner_lists(inner):
                yield x
    else:
        yield lst

>>> lst = [[40, 20, 45, 40], [[20, 10, 10, 30, 45, 20], [30, 20, 20, 30]]]
>>> list(inner_lists(lst))
[[40, 20, 45, 40], [20, 10, 10, 30, 45, 20], [30, 20, 20, 30]]

Or using yield from (Python 3.3 and later):

def inner_lists(lst):
    if all(isinstance(x, list) for x in lst):
        yield from (x for inner in lst for x in inner_lists(inner))
    else:
        yield lst

Or without generators:

def inner_lists(lst):
    if all(isinstance(x, list) for x in lst):
        return [x for inner in lst for x in inner_lists(inner)]
    else:
        return [lst]

Answer 2

If your lists can only contains numbers or sub lists (not both of them) you can use following generator function:

In [25]: def partial_flattening(lst):
             for sub in lst:
                 if isinstance(sub[0], list):
                     yield from sub # in python 2.x use for i in sub: yeild i
                 else:          
                     yield sub
   ....:                 

In [26]: list(partial_flattening(lst))
Out[36]: [[40, 20, 45, 40], [20, 10, 10, 30, 45, 20], [30, 20, 20, 30]]

If there may be combination of both types then based on what you want to do with the numbers you can customize the function.

Or as a more concise approach you can use yield within list comprehension:

In [29]: list([(yield from i) if isinstance(i[0], list) else (yield i) for i in lst])
Out[29]: [[40, 20, 45, 40], [20, 10, 10, 30, 45, 20], [30, 20, 20, 30]]