Comprehensions

Comprehensions

Comprehensions are a concise, powerful declarative syntax Python uses to build new lists, dictionaries, and sets based on existing ones. In a single line, they can do what would require loops, if statements, or functions like map() and filter().

Basic Comprehensions

In their most basic form, comprehensions use the for and in keywords to create instructions that say, “Make a new collection by doing something to every item in this existing collection.” This form of comprehension is like a map() function but more readable.

Here are some examples:

# Create a new list [2, 4, 6, 8]
doubles_list = [number * 2 for number in [1, 2, 3, 4]]

# Create a new set {3, 6, 9, 12}
numbers = [1, 2, 3, 4]
triples_set = {number * 3 for number in numbers}

# Create a new dictionary with
# an additional $10 in each account
accounts = {"savings": 100, "checking": 200}
updated_account = {f"{key}+bonus": accounts[key] + 10 for key in accounts}
# The result:
# {'savings+bonus': 110, 'checking+bonus': 210}

# Create a list with the same items
# as those in a set
condiments = [condiment for condiment in {"ketchup", "mustard", "relish"}]

Filtering Comprehensions With if

If you add one or more if expressions after the for…in expression, you can define criteria to include or exclude items in the collection you’re creating:

Some examples:

# Here's a comprehension acting like a filter() function.
# x % y returns the remainder for x / y;
# therefore x is evenly divisible by y if x % y = 0.
even_numbers = [number for number in range(100) if number % 2 == 0]

# If you use multiple if statements all of them must evaluate to True
# in order to include the item in the new collection.
# '**' is Python's exponent operator.
even_numbers_with_small_squares = [number
 for number in range(100)
  if number % 2 == 0 if number ** 2 < 1000]

# The result:
# [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30]

# You can format the above line
# to be more readable
even_numbers_with_small_squares = [
    number for number in range(100)
    if number % 2 == 0
    if number ** 2 < 1000
]

Conditional Comprehensions With if and else

If you add an if…else expression before the for…in expression, you can specify the value to be added to the new collection if a condition is met and an alternate value if the condition isn’t met.

# Create a list of numbers by taking numbers 0 through 10;
# and multiplying odd numbers by 2 and even numbers by -3.
new_numbers = [
    number * 2 if number % 2 == 1
    else number * -3
    for number in range(11)
]
# The result:
# [0, 2, -6, 6, -12, 10, -18, 14, -24, 18, -30]

# Create a list of food with the correct article
# (food beginning with a vowel starts with "An",
# otherwise it start with "A".)
food_set = {"artichoke", "banana", "coconut", "donut", "egg"}
food_list = [f"An {food}" if food[0] in "aeiou" else f"A {food}" for food
  in food_set]

# The result:
# ['A banana', 'An egg', 'A coconut', 'An artichoke', 'A donut']

Generators

To understand generators, you need to understand iterables and iterators.

An iterable is an object that can return its elements one at a time. You’ve already seen some of Python’s iterables: strings, lists, tuples, dictionaries, and sets. A good general guide is that it’s an iterable if you can access an object’s elements one at a time with a for loop.

An iterator is an object that implements a method named __next__(), which returns the next item in an iterable each time you call it. When there are no more items to return, calling __next__() raises a special exception called StopIteration.

Here’s a quick demonstration of an iterable and its iterator:

planets = ["Mercury", "Venus", "Earth"]

# Create an iterator for `planets`
planet_iterator = iter(planets)

print(next(planet_iterator))  # Mercury
print(next(planet_iterator))  # Venus
print(next(planet_iterator))  # Earth
print(next(planet_iterator))  # StopIteration exception

The code above creates planet_iterator as an iterator for the iterable planets. It uses the built-in next() function to call the iterable’s __next__() method, which returns the next element from planets.

Once planet_iterator has gone through all the elements in planets, trying to iterate once more with the fourth call to next(planet_iterator) raises a StopIteration exception.

Note: If you’re feeling experimental, you can change next(planet_iterator) in the code above to planet_iterator.__next__().

Unless you’re writing very specialized code, you probably won’t use iterators directly. You’re more likely to use for loops, which create their own hidden iterator and quietly handle the StopIteration exception behind the scenes.

Generators

Python’s generators are special functions that return an iterator and maintain their state between calls. They’re another feature that’s better to show than to describe.

Enter the following into a new code cell and run it:

def my_first_generator():
    yield "Here's the first one."
    yield "This would be the second time."
    yield "And now, the third iteration!"

for phrase in my_first_generator():
    print(phrase)

The presence of the keyword yield turns my_first_generator() from an ordinary function into a generator. yield is like return in that it produces a return value and exits the function, but the function is paused rather than terminated, meaning:

  • Any variables in the generator continue to hold their values.
  • When called again by its iterator, the generator will resume from the point after the yield that caused the exit.

The for loop in the code above creates an iterator for my_first_generator() and calls its __next__() method until it raises the StopIteration exception.

Here’s what happens in each iteration:

  • In the first iteration, my_first_generator() exits at the first yield statement, which returns the value "Here's the first one.". The generator is paused, and when its iterator calls it next, it will resume execution on the next line.
  • In the second iteration, my_first_generator() resumes where it left off, the next line, yield "This would be the second time.", which returns its value, pauses the generator and exits.
  • The third iteration resumes my_first_generator() at yield "And now, the third iteration!", once again returning that value, and pausing and exiting the generator.
  • The fourth iteration happens behind the scenes. You don’t experience the halting and error message that goes along with StopIteration. Instead, for handles the exception by terminating the loop.

Generator Comprehensions

In addition to list, set, and dictionary comprehensions, Python supports generator comprehensions, which are delimited by parentheses:

# Generator for the squares of the numbers
# from 0 through 49,999
squared_numbers = (number ** 2 for number in range(50_000))

total_squared_numbers = 0
for item in squared_numbers:
    total_squared_numbers += item
print(total_squared_numbers) # 41665416675000

# Generator for the squares of the *even* numbers
# from 0 through 49,999
squared_even_numbers = (number ** 2 for number in range(50_000) if number % 2 == 0)

# Generator for the squares of the *even* numbers
# and cubes of the *odd* numbers
# from 0 through 49,999
squared_even_cubed_odd_numbers = (
    number ** 2 if number % 2 == 0
    else number ** 3
    for number in range(50_000)
)

Using Generators for Memory Efficiency

If you need a large collection of values to iterate through, you may want to use a generator comprehension instead of a list comprehension. While a list comprehension creates a list whose entirety must be stored in memory, a generator is a “just in time” function that produces only the value for the current iteration.

The following code shows how significant the memory savings from a generator can be:

# Python’s `sys` module contains the getsizeof() function
# which reports the size of an object in bytes
from sys import getsizeof

# Get the size of a list of the squares of
# the first 100 million numbers, starting with 0
getsizeof([number ** 2 for number in range(100_000_000)])
# This should be around 835 million bytes.

# Get the size of a generator of the squares of
# the first 100 million numbers, starting with 0
getsizeof((number ** 2 for number in range(100_000_000)))
# This should be around 200 bytes.

The zip() Function

zip() is a useful built-in function that combines two iterables into a single iterator of tuples. Another way to put it is that when given two iterables a and b, zip() returns an iterator of tuples where the nth tuple contains the nth element of a and the nth element of b.

Having said that, it’s so much easier to demonstrate than explain. Run the following in a new code cell:

numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
months = ("January", "February", "March", "April", "May", "June", "July", "August",
  "September", "October", "November", "December")
birthstones = ("garnet", "amethyst", "aquamarine", "diamond", "emerald",
  "alexandrite", "ruby", "peridot", "sapphire", "tourmaline", "topaz", "tanzanite")
days_of_week = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday",
  "Saturday", "Sunday"]

numbers_and_months = zip(numbers, months)
for number_and_month in numbers_and_months:
    print(number_and_month)

Here’s what the output should look like:

(1, 'January')
(2, 'February')
(3, 'March')
(4, 'April')
(5, 'May')
(6, 'June')
(7, 'July')
(8, 'August')
(9, 'September')
(10, 'October')
(11, 'November')
(12, 'December')

You can zip() more than two iterables:

numbers_months_and_birthstones = zip(numbers, months, birthstones)
for numbers_months_and_birthstone in numbers_months_and_birthstones:
    print(numbers_months_and_birthstone)

Here’s the output of the code above:

(1, 'January', 'garnet')
(2, 'February', 'amethyst')
(3, 'March', 'aquamarine')
(4, 'April', 'diamond')
(5, 'May', 'emerald')
(6, 'June', 'alexandrite')
(7, 'July', 'ruby')
(8, 'August', 'peridot')
(9, 'September', 'sapphire')
(10, 'October', 'tourmaline')
(11, 'November', 'topaz')
(12, 'December', 'tanzanite')

So far, the examples have shown zip()ping iterables of equal lengths. When iterables of different lengths are zip()ped together, zip() stops creating tuples when the shortest iterable is used up:

numbers_and_days_of_week = zip(numbers, days_of_week)
for number_and_day_of_week in numbers_and_days_of_week:
    print(number_and_day_of_week)

The result of zip() is an iterator, but it’s easy to convert it into iterables:

  • List: list(zip(numbers, months))
  • Set: set(zip(numbers, months))
  • Tuple: tuple(zip(numbers, months))

You can also zip() two iterables to make a dictionary. The first iterable becomes the keys, and the second becomes the values. Run the following in a new code cell:

dict(zip(numbers, months))

You should see the following output:

{1: 'January',
 2: 'February',
 3: 'March',
 4: 'April',
 5: 'May',
 6: 'June',
 7: 'July',
 8: 'August',
 9: 'September',
 10: 'October',
 11: 'November',
 12: 'December'}
See forum comments
Download course materials from Github
Previous: Sets Next: Functions: Deep Dive