Comprehensions

A large share of everyday programming is turning one collection into another: the names of the coffees that are low on stock, the prices with tax added, the sales grouped by SKU, the set of customers who bought this week. In PYTH-101, Jordan wrote such transformations as loops that start with an empty list and append to it, and met the list comprehension as a shorter alternative. Experienced Python programmers use comprehensions constantly, not to save lines but because a well-written comprehension states what is being built more directly than the loop that builds it.

This lesson treats comprehensions properly: list comprehensions with filters and conditional expressions, nested comprehensions, dictionary and set comprehensions, and generator expressions passed straight to functions such as sum, any, all and max. It also covers the limits: when a comprehension becomes harder to read than a loop, and how to recognize that point.

List comprehensions

A list comprehension builds a new list from an iterable in a single expression: [expression for item in iterable]. Python runs the for part, evaluates the expression once for each item, and collects the results in order. The loop and the comprehension below produce exactly the same list.

prices = [18.50, 22.00, 19.75, 20.25, 19.00, 21.50]

with_tax = []
for p in prices:
    with_tax.append(round(p * 1.05, 2))
print(with_tax)

with_tax = [round(p * 1.05, 2) for p in prices]
print(with_tax)

[19.43, 23.1, 20.74, 21.26, 19.95, 22.57]
[19.43, 23.1, 20.74, 21.26, 19.95, 22.57]

Read a comprehension from the middle outward: for p in prices says where the items come from, and the expression in front says what to make from each one. The variable p exists only inside the comprehension; it does not leak out into the surrounding code the way a loop variable does. The comprehension is also somewhat faster than the equivalent loop, because Python does not have to look up and call the append method for every item.

Filtering with if

Adding an if clause at the end keeps only the items for which the condition is true. This is the most common shape of all: pick out the items you want and transform them in one step.

stock = {"RC-101": 42, "RC-102": 15, "RC-103": 27, "RC-104": 9, "RC-105": 12, "RC-106": 30}

low = [code for code, bags in stock.items() if bags < 15]
print(low)

big_sellers = [code[-3:] for code, bags in stock.items() if bags >= 27 and code != "RC-106"]
print(big_sellers)

['RC-104', 'RC-105']
['101', '103']

A different construction puts a conditional expression in the expression part, before for. It does not filter: every item produces a result, and the condition chooses which result. Beginners often mix the two up. If you want fewer items, put if at the end; if you want the same number of items with different values, use a conditional expression at the front.

stock = {"RC-101": 42, "RC-104": 9, "RC-105": 12, "RC-106": 30}

flags = ["LOW" if bags < 10 else "ok" for bags in stock.values()]
print(flags)                       # one result per item

only_low = [code for code, bags in stock.items() if bags < 10]
print(only_low)                    # filtered: fewer items

['ok', 'LOW', 'ok', 'ok']
['RC-104']

Nested comprehensions

A comprehension can contain more than one for clause. They are read in the same order as nested loops, left to right being outside to inside. Two common uses are flattening a list of lists into one list, and building every combination of two collections.

orders = [["latte", "muffin"], ["drip"], ["mocha", "latte", "tea"]]
items = [item for order in orders for item in order]
print(items)

sizes, drinks = ["small", "large"], ["latte", "mocha"]
menu = [f"{s} {d}" for d in drinks for s in sizes]
print(menu)

['latte', 'muffin', 'drip', 'mocha', 'latte', 'tea']
['small latte', 'large latte', 'small mocha', 'large mocha']

A comprehension can also contain another comprehension as its expression, which builds a list of lists, such as a table of rows. The row-by-column grid below shows the weekly roast schedule: for each day, the batches planned for each roaster slot.

days = ["Mon", "Tue", "Wed"]
slots = [1, 2, 3]
schedule = [[f"{day}-{slot}" for slot in slots] for day in days]
for row in schedule:
    print(row)

columns = [[row[i] for row in schedule] for i in range(len(slots))]
print(columns[0])

['Mon-1', 'Mon-2', 'Mon-3']
['Tue-1', 'Tue-2', 'Tue-3']
['Wed-1', 'Wed-2', 'Wed-3']
['Mon-1', 'Tue-1', 'Wed-1']

Nested comprehensions are where readability starts to suffer. Two for clauses are usually fine; three, or two with conditions on both, are usually clearer as ordinary loops. The question to ask is not whether something can be written as a comprehension, but whether a reader will understand it faster that way.

def low_names(coffees, limit=10):
    return [c for c in coffees]

def price_list(coffees, rate):
    return [c[2] for c in coffees]

print(low_names(COFFEES))

Dictionary comprehensions

Curly braces with a key: value pair in front build a dictionary instead of a list: {key_expr: value_expr for item in iterable}. Dictionary comprehensions are ideal for building lookups, inverting a dictionary, or transforming every value while keeping the keys.

coffees = [("RC-101", "Sunrise Blend", 18.50), ("RC-102", "Ethiopia Guji", 22.00),
           ("RC-104", "Sumatra Dark", 20.25)]

price_of = {sku: price for sku, name, price in coffees}
print(price_of)

sku_of = {name.lower(): sku for sku, name, price in coffees}
print(sku_of["ethiopia guji"])

raised = {sku: round(p * 1.04, 2) for sku, p in price_of.items() if p < 21}
print(raised)

{'RC-101': 18.5, 'RC-102': 22.0, 'RC-104': 20.25}
RC-102
{'RC-101': 19.24, 'RC-104': 21.06}

If two items produce the same key, the later one wins, just as when assigning to a dictionary in a loop. When inverting a dictionary, that means values that are not unique are silently lost, which is worth checking for. zip is often paired with a dictionary comprehension, or with dict() directly, to combine two parallel lists into one dictionary: dict(zip(codes, stocks)).

Set comprehensions

Curly braces with a single expression build a set: {expression for item in iterable}. Since a set keeps only unique values, a set comprehension is the natural way to ask which distinct things appear in some data, such as the different drinks ordered today or the customers who placed an order this week.

orders = [("Priya", "latte"), ("Sam", "drip"), ("Priya", "muffin"), ("Lee", "latte"), ("Sam", "latte")]

customers = {name for name, item in orders}
latte_fans = {name for name, item in orders if item == "latte"}
print(sorted(customers))
print(sorted(latte_fans))
print(sorted(customers - latte_fans))

['Lee', 'Priya', 'Sam']
['Lee', 'Priya', 'Sam']
[]

Note that empty curly braces, {}, make an empty dictionary, not an empty set; an empty set is written set(). The results here are printed with sorted, because a set has no fixed order and printing it directly could show the names in a different order from one run to the next.

Generator expressions

Round brackets instead of square ones make a generator expression, which produces the items lazily, as in lesson 7, instead of building a list. When a generator expression is the only argument to a function, the extra brackets can be dropped. Generator expressions are the right choice whenever the results are consumed once by a function such as sum, min, max, any, all or str.join, because no intermediate list is needed.

sales = [("RC-101", 2, 18.50), ("RC-104", 1, 20.25), ("RC-102", 3, 22.00), ("RC-101", 1, 18.50)]

print(round(sum(qty * price for sku, qty, price in sales), 2))
print(max(qty for sku, qty, price in sales))
print(any(qty > 2 for sku, qty, price in sales))
print(all(price < 25 for sku, qty, price in sales))
print(", ".join(sorted({sku for sku, qty, price in sales})))

141.75
3
True
True
RC-101, RC-102, RC-104

any and all stop as soon as the answer is known: any at the first true item, all at the first false one. Combined with a generator expression, that means they look at as few items as possible. A common pattern is if any(c.is_low() for c in coffees): print("reorder needed"), which reads almost like the sentence it implements.

Laying out longer comprehensions

A comprehension does not have to fit on one line. Inside brackets, Python ignores line breaks, so a longer comprehension can be split with the expression, the for clause and the if clause each on its own line. This layout, used by code formatters such as Black, makes the three parts easy to see at a glance and is often the difference between a comprehension that reads well and one that should have been a loop.

coffees = [("RC-101", "Sunrise Blend", 18.50, 42), ("RC-102", "Ethiopia Guji", 22.00, 15),
           ("RC-104", "Sumatra Dark", 20.25, 9), ("RC-105", "Decaf Swiss Water", 19.00, 12)]

reorder_lines = [
    f"{sku}  {name:<18} order {30 - stock} bags"
    for sku, name, price, stock in coffees
    if stock < 15
]
print("\n".join(reorder_lines))

RC-104  Sumatra Dark       order 21 bags
RC-105  Decaf Swiss Water  order 18 bags

Python 3.8 added the assignment expression, written with the := operator and nicknamed the walrus operator. It assigns a value and returns it in the same expression, which lets a comprehension compute something once, test it and keep it. In [t for raw in readings if (t := float(raw)) > 190], each reading is converted once, compared, and kept if hot enough. Used sparingly it avoids repeated work; used heavily it makes code dense, so many teams reserve it for cases like this one.

Comprehensions or loops?

Comprehensions are expressions that build a collection. They are the wrong tool when the work is about side effects, such as printing, writing files or changing other objects. A list comprehension used only to call print builds a pointless list of None values. They are also the wrong tool when the logic needs several steps, a try statement, or intermediate variables that deserve names.

raw = ["18.50", "22", "n/a", "19.75", ""]

def to_price(text):
    """Return text as a float, or None if it is not a number."""
    try:
        return float(text)
    except ValueError:
        return None

prices = [p for p in (to_price(t) for t in raw) if p is not None]
print(prices)

# A side effect belongs in a loop, not a comprehension:
for p in prices:
    print(f"{p:6.2f}")

[18.5, 22.0, 19.75]
 18.50
 22.00
 19.75

A useful technique when the logic grows is to move it into a small, named function and call that function from the comprehension, as to_price does here. The function can contain the try statement and be tested on its own, and the comprehension stays a one-line description of the result. The rule of thumb used by many teams: if a comprehension does not fit comfortably on one or two lines, or needs a comment to explain it, write a loop.

def buyers(sales, sku):
    return [c for c, s, q in sales]

def stock_after(stock, sold):
    return stock

print(buyers(SALES, 'RC-101'))