List comprehensions, iterators, generators and generator expressions in Python 3

A list comprehension is a concise way to create lists that would normally require for loops to build.

Example:

list1 = [x**2 for x in range(10)]
print(list1)
#output
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

List comprehension to create a list of tuples:

list2 = [(x, y) for x in [1,2,3] for y in [3,1,4] if x != y]
print(list2)
#output
[(1, 3), (1, 4), (2, 3), (2, 1), (2, 4), (3, 1), (3, 4)]

List comprehension using an if condition:

list3 = [x for x in range(10) if x%2 != 0]
print(list3)
#output
[1, 3, 5, 7, 9]

Only odd numbers are printed in the above example.

Nested list comprehensions allow us to emulate nested for loops in some way.

Example:

matrix = [
    [1,2,3],
    [4,5,6],
    [7,8,9],
    [10,11,12]
]

transposed = [[row[i] for row in matrix] for i in range(3)]
print(transposed)
#output
[[1, 4, 7, 10], [2, 5, 8, 11], [3, 6, 9, 12]]

The equivalent using nested for loops would be:

for i in range(3):
    transposed.append([row[i] for row in matrix])

Moving on to iterators. Iterators let us iterate over container objects using for loops. How to create an iterator: create a class which defines __iter__() and __next__(). __iter__() returns an object with a __next__() method. The __next__() method which is used to retrieve the next object in the container. The __next__() method will also need to raise a StopIteration exception when there are no more elements to iterate over.

Example:

class Squared:
    """Square all the numbers"""

    def __init__(self, data):
        self.data = data
        self.index = 0
    
    def __iter__(self):
        return self

    def __next__(self):
        if self.index == len(self.data):
            raise StopIteration
        val = self.data[self.index] **2
        self.index += 1

        return val

values = Squared([1,2,3,10])

for i in values:
    print(i)

#output
1
4
9
100

Generators are a way to create iterators. There is no need to explicitly define the __iter__() and __next__() methods, they are created automatically. Generators are functions that use the yield statement to return data. When __next__() is called, the function resumes where it left off and remembers the state of the program. The StopIteration exception is also automatically raised.

def double_values(data):
    """Doubles all the values"""
    for val in data:
        yield val*2

for i in double_values([4,5,3]):
    print(i)

#output
8
10
6

Generator expressions are a simple but limited way to create generators and used in cases where the return value of the generator expression is used immediately. They have a syntax similar to list comprehensions, but use parantheses instead of brackets.

Since generator expressions generate the values on the fly instead of storing all the values in memory like list comprehensions do, they tend to be more memory efficient than the equivalent list comprehension but also tend to be a little slower. This is an important tradeoff to keep in mind.

Example:

exp = sum(i for i in range(5))
print(exp)

#output
10

Source code for today’s plog is here.

References: