Search code examples
pythonpython-3.xlist-comprehensiongenerator-expression

Generator expression in list comprehension not working as expected


The following code produces expected output:

# using a list comprehension as the first expression to a list comprehension
>>> l = [[i*2+x for x in j] for i,j in zip([0,1],[range(4),range(4)])]
>>> l[0]
[0, 1, 2, 3]
>>> l[1]
[2, 3, 4, 5]

However, when I use a generator expression instead, I get a different result:

# using a generator expression as the first expression
>>> l = [(i*2+x for x in j) for i,j in zip([0,1],[range(4),range(4)])]
>>> list(l[0])
[2, 3, 4, 5]
>>> list(l[1])
[2, 3, 4, 5]
>>> list(l[0])
[]
>>> list(l[1])
[]
>>> l
[<generator object <listcomp>.<genexpr> at 0x7fddfa413ca8>, <generator object <listcomp>.<genexpr> at 0x7fddfa413c50>]

I understand that generator expressions can only be used once, however I'm having trouble reasoning why I'm getting the same list twice in this scenario, especially since the generator objects appear to be unique.

What am I missing here? This was tested on Python 3.6.5.


Solution

  • The generator objects are unique, but they refer to i and j, but the the list-comprehension terminates (which essentially creates a function scope, just like the generator expressions inside the list-comprehension). Thus, i and j have the values i == 1 and j == range(4). You can even introspect this:

    In [1]: l = [(i*2+x for x in j) for i,j in zip([0,1],[range(4),range(4)])]
    
    In [2]: g = l[0]
    
    In [3]: g.gi_frame.f_locals
    Out[3]: {'.0': <range_iterator at 0x10e9be960>, 'i': 1}
    

    This is essentially the same reason why this often surprising behavior occurs:

    In [4]: fs = [lambda: i for i in range(3)]
    
    In [5]: fs[0]
    Out[5]: <function __main__.<listcomp>.<lambda>()>
    
    In [6]: fs[0]()
    Out[6]: 2
    
    In [7]: fs[1]()
    Out[7]: 2
    
    In [8]: fs[2]()
    Out[8]: 2
    

    You can fix this using the same solution, which is to create another enclosing scope, which binds the variables locally to something that won't change. Using a function (a lambda here, but it could be regular function) would work perfectly:

    In [9]: l = [(lambda i, j: (i*2+x for x in j))(i, j) for i,j in zip([0,1],[range(4),range(4)])]
    
    In [10]: list(l[0])
    Out[10]: [0, 1, 2, 3]
    
    In [11]: list(l[1])
    Out[11]: [2, 3, 4, 5]
    

    Although, perhaps for clarity, I'll use different parameter names to make it more obvious what is going on:

    In [12]: l = [(lambda a, b: (a*2+x for x in b))(i, j) for i,j in zip([0,1],[range(4),range(4)])]
    
    In [13]: list(l[0])
    Out[13]: [0, 1, 2, 3]
    
    In [14]: list(l[1])
    Out[14]: [2, 3, 4, 5]