Python Comprehensions
A comprehension is a short way to build a new list, dictionary or set from another collection.
From Loop to Comprehension
You met list comprehensions briefly in the Lists lesson. Here is the full story. First, the normal loop way:
Example
prices = [120, 450, 80, 1000]
doubled = []
for p in prices:
doubled.append(p * 2)
print(doubled)[240, 900, 160, 2000]
The same thing as a comprehension:
Example
prices = [120, 450, 80, 1000]
doubled = [p * 2 for p in prices]
print(doubled)[240, 900, 160, 2000]
The pattern is always the same:
Example
new_list = [expression for item in collection]Read it from the middle: "for every item in collection, put expression in the new list."
Adding a Filter
Put an if at the end to keep only some items.
Example
marks = {"Ali": 72, "Sara": 45, "Bilal": 88, "Hina": 39}
passed = [name for name in marks if marks[name] >= 50]
print(passed)
words = ["karachi", "", "lahore", "", "quetta"]
cleaned = [w.title() for w in words if w]
print(cleaned)['Ali', 'Bilal'] ['Karachi', 'Lahore', 'Quetta']
An empty string is falsy, so if w skips the blanks.
if...else Inside
To change every item in one of two ways, put if...else before the for:
Example
marks = [72, 45, 88, 39]
results = ["Pass" if m >= 50 else "Fail" for m in marks]
print(results)['Pass', 'Fail', 'Pass', 'Fail']
Where the if goes | What it does |
|---|---|
[x for x in data if x > 0] | Filters: some items are dropped |
[a if x > 0 else b for x in data] | Chooses: every item is kept, but changed |
Dictionary Comprehensions
Use curly brackets and a key: value pair.
Example
items = ["tea", "samosa", "juice"]
prices = [60, 40, 120]
menu = {item: price for item, price in zip(items, prices)}
print(menu)
with_tax = {k: round(v * 1.18) for k, v in menu.items()}
print(with_tax)
cheap = {k: v for k, v in menu.items() if v < 100}
print(cheap){'tea': 60, 'samosa': 40, 'juice': 120}
{'tea': 71, 'samosa': 47, 'juice': 142}
{'tea': 60, 'samosa': 40}zip() pairs the two lists item by item, which is perfect for building a dictionary.
Set Comprehensions
Curly brackets without a colon make a set, so duplicates disappear.
Example
emails = ["ali@gmail.com", "sara@yahoo.com", "omar@gmail.com", "hina@outlook.com"]
domains = {e.split("@")[1] for e in emails}
print(sorted(domains))['gmail.com', 'outlook.com', 'yahoo.com']
Nested Loops
A comprehension can have two for parts. They run in the same order as nested loops.
Example
sizes = ["S", "M"]
colours = ["red", "black"]
shirts = [f"{c}-{s}" for c in colours for s in sizes]
print(shirts)
table = [[1, 2, 3], [4, 5, 6]]
flat = [n for row in table for n in row]
print(flat)['red-S', 'red-M', 'black-S', 'black-M'] [1, 2, 3, 4, 5, 6]
Generator Expressions
Use round brackets instead of square ones and you get a generator expression. It makes values one at a time instead of building a whole list, which saves memory. It is often passed straight into sum(), max() or any().
Example
bill = [("tea", 60, 3), ("samosa", 40, 5)]
total = sum(price * qty for name, price, qty in bill)
print(total)380
You will learn more about these in Iterators and Generators.
Comprehensions are great for short, simple jobs. If yours needs several conditions or does not fit on one line, a normal for loop is easier to read. Clear code beats clever code.
Exercise
Exercise
What does [n * n for n in range(1, 6) if n % 2 == 1] give?