Comprehensions
Many loops do the same job: start with an empty list, go through some items, and append something for each one. Python has a short way to write exactly that, called a comprehension. Here is the long way first:
1loads = [3, 8, 2, 6]2doubled = []3for ore in loads:4 doubled.append(ore * 2)5print(doubled)
[6, 16, 4, 12]
From a loop to one line
The same list, built with a list comprehension:
1loads = [3, 8, 2, 6]2doubled = [ore * 2 for ore in loads]3print(doubled)
[6, 16, 4, 12]
Read it as "ore * 2, for each ore in loads". The pieces of the loop are all still there, just moved around:
| Piece | In the loop | In the comprehension |
|---|---|---|
| what to collect | doubled.append(ore * 2) | ore * 2 at the front |
| where items come from | for ore in loads: | for ore in loads after it |
| the new list | doubled = [] before the loop | the square brackets around it all |
Tip: Press Step on both examples. The loop goes around line by line; the comprehension also takes one step per item, building the list as it goes.
Changing every item
The part at the front can be any expression: a calculation, a string method, an f-string. The for part can loop over anything a for loop can, range() included.
1names = ["ada", "ben", "cleo"]2print([name.title() for name in names])3print([len(name) for name in names])4print([n * n for n in range(1, 6)])5print([f"lamp {n}" for n in range(3)])
['Ada', 'Ben', 'Cleo'] [3, 3, 4] [1, 4, 9, 16, 25] ['lamp 0', 'lamp 1', 'lamp 2']
Keeping only some items
Add an if at the end, and only the items that pass the test go into the new list. The others are skipped.
1loads = [3, 8, 2, 6, 11]2big = [ore for ore in loads if ore >= 5]3print(big)4words = ["mine", "ore", "lamp", "warehouse"]5print([w.upper() for w in words if len(w) > 3])
[8, 6, 11] ['MINE', 'LAMP', 'WAREHOUSE']
What does this print?
1numbers = [1, 2, 3, 4, 5, 6]2print([n * 10 for n in numbers if n % 2 == 0])
The if keeps only the even numbers, 2, 4 and 6. Each of those then goes through n * 10.
Choosing a value with if-else
To keep every item but choose what to put in for each one, use a conditional expression, a if test else b, at the front:
1loads = [3, 8, 2, 6]2sizes = ["big" if ore >= 5 else "small" for ore in loads]3print(sizes)
['small', 'big', 'small', 'big']
The position of the if changes what it does:
| Where | Written | Effect |
|---|---|---|
| at the end | [x for x in items if test] | filters: the list can get shorter |
| at the front | [a if test else b for x in items] | chooses: the list stays the same length |
How many numbers are in the printed list?
1print([n if n > 0 else 0 for n in [4, -2, 7, -5]]) An if-else at the front never drops an item; it only chooses the value. The result is [4, 0, 7, 0]: still four numbers.
Dictionary comprehensions
Use curly braces and a key: value pair at the front, and you build a dictionary instead of a list. zip() and .items() work here just as in a for loop.
1names = ["ada", "ben", "cleo"]2print({name: len(name) for name in names})3prices = {"iron": 4, "copper": 6, "gold": 20}4print({item: price * 2 for item, price in prices.items()})5print({item: price for item, price in prices.items() if price > 5})
{'ada': 3, 'ben': 3, 'cleo': 4}
{'iron': 8, 'copper': 12, 'gold': 40}
{'copper': 6, 'gold': 20}
Set comprehensions
Curly braces with a single value at the front build a set, so duplicates disappear. Here, which remainders do we get when we divide 0 to 9 by 4?
1print({n % 4 for n in range(10)})2words = ["Lamp", "lamp", "ORE", "ore", "Mine"]3print(sorted({w.lower() for w in words}))
{0, 1, 2, 3}
['lamp', 'mine', 'ore']
A set of strings can print in a different order each time a program runs, so the second line wraps it in sorted() to get a list in a fixed order.
| Written as | Front part | You get |
|---|---|---|
[ ... ] | a value | a list |
{ ... } | key: value | a dictionary |
{ ... } | a value | a set |
When not to use one
A comprehension is for building one new collection. Keep a normal for loop when:
- the loop does something, like printing, instead of collecting values;
- you need
break, or several steps for each item; - the comprehension gets so long that you have to read it twice.
This loop prints, stops early and keeps a total, so it should stay a loop:
1loads = [4, 7, 0, 5]2total = 03for ore in loads:4 if ore == 0:5 print("Empty cart, stop")6 break7 total += ore8 print(f"Unloaded {ore}, total {total}")
Unloaded 4, total 4 Unloaded 7, total 11 Empty cart, stop
Short and clear wins. If a comprehension is not clearer than the loop, write the loop.
Practice
Write the code, then press Check. A hint costs nothing; the answer is there when you need it. Each one fits on a single line.
Even squares
✓Replace the empty list with a comprehension that holds the square of every even number in numbers.
[4, 16, 36, 64, 100]
1numbers = list(range(1, 11))2squares = []3print(squares)
Lit or dark
✓Turn the list of oil levels into a list of words: "lit" when the level is above 0, otherwise "dark".
['lit', 'dark', 'lit', 'dark', 'lit']
1oil = [5, 0, 2, 0, 7]2states = [level for level in oil]3print(states)
Restock order
✓Build a dictionary of what to order: for each item with fewer than 10 in stock, how many more are needed to reach 10.
{'copper': 7, 'gold': 10}1stock = {"iron": 14, "copper": 3, "coal": 22, "gold": 0}2order = {}3print(order)
What you learned
[expression for x in items]builds a new list in one line, instead of an empty list plusappendin a loop.- An
ifat the end filters items out;a if test else bat the front chooses a value for every item. {key: value for ...}builds a dictionary;{value for ...}builds a set.- Comprehensions build collections. For printing,
breakor several steps per item, use a normal loop. - If a comprehension is hard to read, the loop is the better choice.