List, dict, set comprehension in Python: concise yet readable

1. Overview

Comprehension is a distinctive syntax in Python. It allows you to create a list, dict, or set from an iterable in a single line. Written correctly, the code is shorter and clearer than a for loop. Overuse turns code into a long, hard-to-understand line.

This article goes through list, dict, set comprehension and generator expressions with concrete examples. After reading, you’ll know when to use comprehension, when to go back to a for loop, and how to write tests for data transformation logic. All examples run on Python 3.12.

2. List comprehension

2.1. Basic syntax

Suppose you have a list of orders and need to extract the total amount of each order. The traditional way to write it using a for loop is as follows:

from dataclasses import dataclass

@dataclass
class Order:
    id: int
    customer: str
    total: float
    status: str

orders = [
    Order(1, "An", 250_000, "PAID"),
    Order(2, "Bình", 90_000, "CANCELLED"),
    Order(3, "An", 1_200_000, "PAID"),
    Order(4, "Chi", 450_000, "PENDING"),
]

totals = []
for order in orders:
    totals.append(order.total)
print(totals)

Rewrite using list comprehension:

totals = [order.total for order in orders]
print(totals)
[250000, 90000, 1200000, 450000]

The structure is [expression for element in iterable]. Read from left to right: “take order.total for each order in orders”. You don’t need to declare an empty list or call append. The intent “I’m creating a new list” is immediately apparent from the [ bracket.

2.2. Filtering elements with if

If you only want to take paid orders, you add an if clause at the end:

paid_ids = [order.id for order in orders if order.status == "PAID"]
print(paid_ids)
[1, 3]

The if at the end acts as a filter. Elements that don’t meet the condition are skipped, so the result list can be shorter than the original list.

2.3. if-else expression at the beginning

When you want to transform every element based on a condition instead of filtering, you place the if-else in the expression part:

labels = ["large" if order.total >= 500_000 else "small" for order in orders]
print(labels)
['small', 'small', 'large', 'small']

This is where many people get confused. You just need to remember two rules:

  • if after for is used to filter. This clause has no else.
  • a if cond else b before for is a conditional expression. This clause must have else and keeps the number of elements the same.

Writing [x for x in data if cond else y] will cause a SyntaxError.

3. Dict comprehension

Dict comprehension uses {} with key-value pairs. For example, you need to quickly look up an order by id:

orders_by_id = {order.id: order for order in orders}
print(orders_by_id[3])
Order(id=3, customer='An', total=1200000, status='PAID')

You can also transform an existing dict. For example, apply a 10% tax to a price table:

prices = {"cafe": 30_000, "tra_sua": 45_000, "banh_mi": 20_000}
prices_with_vat = {name: round(price * 1.1) for name, price in prices.items()}
print(prices_with_vat)
{'cafe': 33000, 'tra_sua': 49500, 'banh_mi': 22000}

prices.items() returns pairs of (key, value), and you unpack them directly in the for loop.

A useful trick is to reverse a dict:

status_codes = {"PAID": 1, "PENDING": 2, "CANCELLED": 3}
code_to_status = {code: status for status, code in status_codes.items()}
print(code_to_status[2])
PENDING

Note: If two elements produce the same key, the later element will silently overwrite the earlier one. For example, {order.customer: order.total for order in orders} only keeps the last order for “An” (1,200,000), and the 250,000 order disappears with no warning. When you need to group, use defaultdict with a for loop (see section 7).

4. Set comprehension

Set comprehension also uses {} but has only one expression, with no :. The result automatically removes duplicates. For example, you want to know which customers have placed orders:

customers = {order.customer for order in orders}
print(sorted(customers))
['An', 'Bình', 'Chi']

Each customer appears only once even though “An” has two orders. The order in a set is not guaranteed, so the example above uses sorted when printing.

Set comprehension is very convenient when you need to quickly check membership:

banned_emails = {"spam@x.com", "Fake@Y.com"}
normalized = {email.strip().lower() for email in banned_emails}
print("fake@y.com" in normalized)
True

Note: Empty {} is a dict, not a set. To create an empty set, you must write set().

5. Multiple for loops and nested comprehensions

You can write multiple for loops in a single comprehension. The reading order is identical to when you write nested for loops from top to bottom:

carts = {
    "An": ["cafe", "banh_mi"],
    "Chi": ["tra_sua"],
}

items = [(customer, item) for customer, products in carts.items() for item in products]
print(items)
[('An', 'cafe'), ('An', 'banh_mi'), ('Chi', 'tra_sua')]

The comprehension above is equivalent to the following code:

items = []
for customer, products in carts.items():
    for item in products:
        items.append((customer, item))

The most common case is flattening a nested list:

matrix = [[1, 2, 3], [4, 5, 6]]
flat = [value for row in matrix for value in row]
print(flat)
[1, 2, 3, 4, 5, 6]

Way to remember: the outermost for loop comes first. If you have to pause and read a second time to understand the order, that’s a sign you should separate it into a regular for loop.

6. Generator expression: when you don’t need a full list

If you replace [] with (), you’ll have a generator expression. Python doesn’t create a list in memory. Values are generated one at a time as code consumes them.

revenue = sum(order.total for order in orders if order.status == "PAID")
print(revenue)
1450000

When a generator is the only argument to a function, you can omit a pair of parentheses. This works well with sum, any, all, max, min:

has_big_order = any(order.total > 1_000_000 for order in orders)
all_positive = all(order.total > 0 for order in orders)
print(has_big_order, all_positive)
True True

any stops as soon as it finds the first element that satisfies the condition. So with large data, a generator both saves memory and is faster than any([...]).

Note: A generator can only be iterated once. A second iteration will return nothing:

paid = (o for o in orders if o.status == "PAID")
print(len(list(paid)))
print(len(list(paid)))
2
0

If you need to use the result multiple times, use list comprehension.

7. Keeping comprehensions readable

7.1. Name functions for complex logic

Long conditions make comprehensions hard to read:

vip = [o.customer for o in orders if o.status == "PAID" and o.total >= 1_000_000 and not o.customer.startswith("test")]

You can separate the condition into a clearly named function:

def is_vip_order(order: Order) -> bool:
    return (
        order.status == "PAID"
        and order.total >= 1_000_000
        and not order.customer.startswith("test")
    )

vip = [order.customer for order in orders if is_vip_order(order)]
print(vip)
['An']

Now the comprehension reads like a sentence, while the business logic has its own place to be tested.

7.2. Strategic line breaks

Comprehensions don’t have to fit on a single line. Within the parentheses, you can break lines freely:

summary = {
    order.id: order.total
    for order in orders
    if order.status != "CANCELLED"
}
print(summary)
{1: 250000, 3: 1200000, 4: 450000}

Each line has a role: expression, source, filter.

7.3. Walrus operator to avoid computing twice

Sometimes you need both to filter and use the result of the same computation. The := operator (Python 3.8+) lets you compute only once:

def discount_for(order: Order) -> float:
    return order.total * 0.05 if order.total >= 400_000 else 0

discounts = {
    order.id: d
    for order in orders
    if (d := discount_for(order)) > 0
}
print(discounts)
{3: 60000.0, 4: 22500.0}

Use walrus sparingly. The variable d will “leak” into the outer scope of the comprehension. Also, many readers aren’t familiar with this syntax yet.

8. Testing data transformation logic

Comprehensions often live in data transformation functions, so you should test them like any other logic. First, install pytest:

pip install pytest

Suppose you have a module reports.py:

# reports.py
from dataclasses import dataclass

@dataclass
class Order:
    id: int
    customer: str
    total: float
    status: str

def revenue_by_customer(orders: list[Order]) -> dict[str, float]:
    customers = {o.customer for o in orders if o.status == "PAID"}
    return {
        c: sum(o.total for o in orders if o.customer == c and o.status == "PAID")
        for c in customers
    }

And the corresponding test file:

# test_reports.py
from reports import Order, revenue_by_customer

def test_revenue_only_counts_paid_orders():
    orders = [
        Order(1, "An", 250_000, "PAID"),
        Order(2, "An", 1_200_000, "PAID"),
        Order(3, "Bình", 90_000, "CANCELLED"),
        Order(4, "Chi", 450_000, "PAID"),
    ]

    result = revenue_by_customer(orders)

    assert result == {"An": 1_450_000, "Chi": 450_000}
    assert "Bình" not in result

def test_empty_orders_returns_empty_dict():
    assert revenue_by_customer([]) == {}
pytest -q
..                                                                       [100%]
2 passed in 0.01s

The first test asserts two things: cancelled orders aren’t counted, and customers with only cancelled orders don’t appear in the result. The second test checks the edge case with an empty list.

The function above works but has O(n × number of customers) complexity because each customer scans the entire orders. With large data, a for loop with defaultdict is both faster and more readable:

from collections import defaultdict

def revenue_by_customer(orders: list[Order]) -> dict[str, float]:
    totals: dict[str, float] = defaultdict(float)
    for order in orders:
        if order.status == "PAID":
            totals[order.customer] += order.total
    return dict(totals)

The old tests still pass without modification. This is the benefit of having tests: you can confidently refactor from comprehension to for loop or vice versa.

9. When NOT to use comprehension

  • When you have side effects. Don’t write [print(o) for o in orders] or [db.save(o) for o in orders]. This creates a meaningless list of Nones and misleads readers into thinking you need a result. Use a for loop.

  • When nested more than two levels deep or with multiple interleaved ifs. Three for loops in one line are almost always harder to read than the multi-line version.

  • When you need grouping or accumulation like summing or counting. Use defaultdict, Counter, or a for loop.

  • When you need to handle exceptions per element. Comprehensions can’t contain try/except. You should separate it into a function and call it in the comprehension, or use a for loop.

  • When data is huge and you only iterate once. Then use a generator expression instead of a list.

Note: In Python 3, the loop variable of a comprehension does not leak out. After [x for x in range(3)], the name x doesn’t exist in the outer scope, unlike a regular for loop. The only exception is a variable assigned with := as mentioned in section 7.3.

10. Conclusion

Comprehensions help you express the idea “create a new collection from an old one” concisely:

  • [...] for list, {k: v ...} for dict, {...} for set, and (...) for generator.
  • if after for filters, while a if cond else b before for transforms.
  • Separate complex conditions into clearly named functions, and break lines when expressions are long.
  • Use generator expressions with sum, any, all to save memory.
  • Return to for loops when you have side effects, need grouping, handle exceptions, or when comprehensions become hard to read.

The simplest rule: if a colleague has to read your comprehension twice to understand it, rewrite it as a for loop. Concise code only has value if it’s still easy to read.

pythoncomprehensionclean codegeneratorpytest