Comprehensions and Iteration

Ava needs a list of every price with tax. The loop she writes takes four lines, and a friend smiles, and shows her a single line that does the same. Python programmers love this style. It is called a comprehension, and once you read one, you will see them everywhere.

• The task: transform a whole list
• The old way: a four-line loop
• The Pythonic way: one clear line

The list comprehension

A list comprehension builds a new list from an old one. Read it like English. Give me this expression, for each item, in this list. The first version is the loop we know, and the second does the same in one line. Both give the same answer, but the second is shorter, and often faster.

prices = [10, 20, 30]
with_tax = []
for p in prices:
    with_tax.append(p * 1.2)
print(with_tax)
new = [p * 1.2 for p in prices]
print(new)

[12.0, 24.0, 36.0]
[12.0, 24.0, 36.0]

Adding a condition

You can add an if at the end, to keep only some items. Here, we keep only the cheap plants, and we also build a list of upper case names. You can even choose between two values with an if else placed before the for. Keep them short. If one gets complicated, use a normal loop instead.

prices = [5, 25, 12, 40, 8]
cheap = [p for p in prices if p < 15]
labels = ["cheap" if p < 15 else "premium" for p in prices]
names = [n.upper() for n in ["fern", "ivy"]]
print(cheap)
print(labels)
print(names)

[5, 12, 8]
['cheap', 'premium', 'cheap', 'premium', 'cheap']
['FERN', 'IVY']

Dictionary and set comprehensions

The same idea builds dictionaries, with a key, a colon and a value, and sets, with just curly brackets. Here Ava creates a price dictionary from two lists, using zip, which pairs them up, and a set of unique first letters. It is a compact way to reshape data.

names = ["fern", "ivy", "rose"]
prices = [9.5, 7.0, 11.0]
menu = {n: p for n, p in zip(names, prices)}
letters = {n[0] for n in names}
print(menu)
print(letters)
print({n: len(n) for n in names})

{'fern': 9.5, 'ivy': 7.0, 'rose': 11.0}
{'i', 'f', 'r'}
{'fern': 4, 'ivy': 3, 'rose': 4}

Handy iteration tools

A few built in tools make loops cleaner. Enumerate gives index and item. Zip walks several lists together. Reversed and sorted change the direction and order. And any and all test a whole list at once, answering, is anything true, or is everything true. Combine them, and your loops get short and clear.

Comprehension or loop?

So which should you choose? Use a comprehension when you build a list from another, simply and clearly. Use a loop when there are many steps, side effects like printing, or complex logic. The rule is readability. If you need to pause to understand a line, make it a loop. Clear code always wins over clever code.

Square the even numbers

Your turn. Using one list comprehension, build a list of the squares of only the even numbers from one to ten. Ready? Here is one solution. The expression is x times x, the loop is over range, and the condition keeps the even numbers.

evens = [x * x for x in range(1, 11) if x % 2 == 0]
print(evens)

[4, 16, 36, 64, 100]

Summary