Search code examples
pythongenerator

How does a generator comprehension works?


What does generator comprehension do? How does it work? I couldn't find a tutorial about it.


Solution

  • Do you understand list comprehensions? If so, a generator expression is like a list comprehension, but instead of finding all the items you're interested and packing them into list, it waits, and yields each item out of the expression, one by one.

    >>> my_list = [1, 3, 5, 9, 2, 6]
    >>> filtered_list = [item for item in my_list if item > 3]
    >>> print(filtered_list)
    [5, 9, 6]
    >>> len(filtered_list)
    3
    >>> # compare to generator expression
    ... 
    >>> filtered_gen = (item for item in my_list if item > 3)
    >>> print(filtered_gen)  # notice it's a generator object
    <generator object <genexpr> at 0x7f2ad75f89e0>
    >>> len(filtered_gen) # So technically, it has no length
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
    TypeError: object of type 'generator' has no len()
    >>> # We extract each item out individually. We'll do it manually first.
    ... 
    >>> next(filtered_gen)
    5
    >>> next(filtered_gen)
    9
    >>> next(filtered_gen)
    6
    >>> next(filtered_gen) # Should be all out of items and give an error
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
    StopIteration
    >>> # Yup, the generator is spent. No values for you!
    ... 
    >>> # Let's prove it gives the same results as our list comprehension
    ... 
    >>> filtered_gen = (item for item in my_list if item > 3)
    >>> gen_to_list = list(filtered_gen)
    >>> print(gen_to_list)
    [5, 9, 6]
    >>> filtered_list == gen_to_list
    True
    >>> 
    

    Because a generator expression only has to yield one item at a time, it can lead to big savings in memory usage. Generator expressions make the most sense in scenarios where you need to take one item at a time, do a lot of calculations based on that item, and then move on to the next item. If you need more than one value, you can also use a generator expression and grab a few at a time. If you need all the values before your program proceeds, use a list comprehension instead.