Python in Detail: Generator Comprehensions

• Generator functions are terse, but sometimes we can replace them with generator comprehensions, which are even shorter. These look exactly like list comprehensions, except that they're surrounded by (parens) instead of [square brackets].

• Like generator functions, generator comprehensions return an iterator. In some of these examples, we'll call list on the iterators to see their contents.

• >

  my_iter = (x for x in range(0, 6) if x % 2 == 0)
  list(my_iter)

  Result:

  Run Code

• Or we can call next on the iterator, as usual.

• >

  my_iter = (char for char in "hello" if char != "h")
  next(my_iter)

  Result:

  'e'

• Those examples used finite data like range(0, 6) and "hello", so we could've used list comprehensions rather than generator comprehensions. However, generator comprehensions also work on infinitely-long iterators.

• >

  import itertools
  
  my_iter = (n for n in itertools.count(3) if n % 3 == 0)
  (next(my_iter), next(my_iter), next(my_iter), next(my_iter))

  Result:

  Run Code

• If we wrote that as a list expression, it would eventually consume all available memory and crash.

• Let's look at one more example to highlight how much code we can save with generator comprehensions. In an earlier lesson, we wrote Primes and PrimesIterator classes to iterate over all prime numbers.

• >

  class Primes:
    def __iter__(self):
      return PrimesIterator()
  
  class PrimesIterator:
    def __init__(self):
      self.current_prime = 2
  
    def __next__(self):
      current_prime = self.current_prime
      self.increment_to_next_prime()
      return current_prime
  
    def increment_to_next_prime(self):
      while True:
        self.current_prime += 1
        if is_prime(self.current_prime):
          return
  
  def is_prime(n):
    for i in range(2, n // 2 + 1):
      if n % i == 0:
        return False
    return True
• Note: this code example reuses elements (variables, etc.) defined in earlier examples.

  >

  my_iter = iter(Primes())
  (next(my_iter), next(my_iter), next(my_iter), next(my_iter), next(my_iter))

  Result:

  (2, 3, 5, 7, 11)

• When we covered generator functions, we simplified that code quite a bit, giving us the version below.

• >

  def primes():
    n = 2
    while True:
      if is_prime(n):
        yield n
      n += 1
  
  def is_prime(n):
    for i in range(2, n // 2 + 1):
      if n % i == 0:
        return False
    return True
  
  my_iter = primes()
  (next(my_iter), next(my_iter), next(my_iter), next(my_iter), next(my_iter))

  Result:

  (2, 3, 5, 7, 11)

• Now we can use generator comprehensions to simplify it even more. We still need an is_prime function to calculate prime numbers. But we can reduce the primes() generator to a single line! We'll use itertools.count(2) to iterate through numbers starting at 2, instead of keeping track of n ourselves.

• >

  import itertools
  
  primes = (n for n in itertools.count(2) if is_prime(n))
  
  def is_prime(n):
    for i in range(2, n // 2 + 1):
      if n % i == 0:
        return False
    return True
  
  (next(primes), next(primes), next(primes), next(primes), next(primes))

  Result:

  (2, 3, 5, 7, 11)

• Other than the is_prime function, all of the other code in the original pair of classes collapsed into a one-line generator comprehension. That's 15 lines of code, in two classes, replaced with a single medium-length line! Generators, and especially generator comprehensions, can be extremely terse!

• However, we should also be careful not to overuse generators and generator comprehensions. It's easy to get carried away, trying to make code as small as possible because it's a fun challenge.

• We could take the example above even farther: it's possible to write is_prime itself using another generator comprehension. That would made the code even shorter, but also less readable. Overusing comprehensions makes code less comprehensible.

• When in doubt, we suggest waiting a day, then re-reading the generator comprehension. If it's not immediately obvious to you after only one day away, then it will certainly confuse other people reading it for the first time!