List, dict, set comprehension trong Python: viết gọn mà vẫn dễ đọc

1. Tổng quan

Comprehension là một cú pháp đặc trưng của Python. Nó cho phép bạn tạo list, dict hoặc set từ một iterable chỉ trong một dòng. Viết đúng cách thì code ngắn và rõ hơn vòng for. Lạm dụng thì code thành một dòng dài khó hiểu.

Bài này đi qua list, dict, set comprehension và generator expression bằng ví dụ cụ thể. Đọc xong, bạn sẽ biết khi nào nên dùng comprehension, khi nào nên quay lại vòng for, và cách viết test cho logic biến đổi dữ liệu. Toàn bộ ví dụ chạy trên Python 3.12.

2. List comprehension

2.1. Cú pháp cơ bản

Giả sử bạn có danh sách đơn hàng và cần lấy tổng tiền của từng đơn. Cách viết bằng vòng for truyền thống như sau:

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)

Viết lại bằng list comprehension:

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

[250000, 90000, 1200000, 450000]

Cấu trúc là [biểu_thức for phần_tử in iterable]. Đọc từ trái sang phải: “lấy order.total cho mỗi order trong orders”. Bạn không cần khai báo list rỗng và cũng không cần gọi append. Ý đồ “tôi đang tạo một list mới” lộ ra ngay từ dấu [.

2.2. Lọc phần tử với if

Nếu chỉ muốn lấy các đơn đã thanh toán, bạn thêm mệnh đề if ở cuối:

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

[1, 3]

if ở cuối đóng vai trò bộ lọc. Phần tử nào không thỏa điều kiện sẽ bị bỏ qua, nên list kết quả có thể ngắn hơn list gốc.

2.3. Biểu thức if-else ở đầu

Khi muốn biến đổi mọi phần tử theo điều kiện thay vì lọc, bạn đặt if-else ở phần biểu thức:

labels = ["lớn" if order.total >= 500_000 else "nhỏ" for order in orders]
print(labels)

['nhỏ', 'nhỏ', 'lớn', 'nhỏ']

Đây là chỗ nhiều người nhầm. Bạn chỉ cần nhớ hai quy tắc:

• if sau for dùng để lọc. Mệnh đề này không có else.
• a if cond else b trước for là biểu thức điều kiện. Mệnh đề này bắt buộc có else và giữ nguyên số phần tử.

Viết [x for x in data if cond else y] sẽ gây SyntaxError.

3. Dict comprehension

Dict comprehension dùng dấu {} với cặp key: value. Ví dụ, bạn cần tra cứu đơn hàng nhanh theo 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')

Bạn cũng có thể biến đổi một dict có sẵn. Ví dụ, áp thuế 10% cho bảng giá:

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() trả về các cặp (key, value), và bạn unpack chúng ngay trong vòng for.

Một mẹo hữu ích là đảo ngược 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

Lưu ý: Nếu hai phần tử cho ra cùng một key, phần tử sau sẽ ghi đè im lặng phần tử trước. Ví dụ {order.customer: order.total for order in orders} chỉ giữ đơn cuối cùng của “An” (1.200.000), và đơn 250.000 biến mất mà không có cảnh báo nào. Khi cần gom nhóm, hãy dùng defaultdict với vòng for (xem mục 7).

4. Set comprehension

Set comprehension cũng dùng {} nhưng chỉ có một biểu thức, không có dấu :. Kết quả tự loại trùng. Ví dụ, bạn muốn biết những khách nào đã đặt hàng:

customers = {order.customer for order in orders}
print(sorted(customers))

['An', 'Bình', 'Chi']

Mỗi khách chỉ xuất hiện một lần dù “An” có hai đơn. Thứ tự trong set không được đảm bảo, nên ví dụ trên dùng sorted khi in.

Set comprehension rất tiện khi cần kiểm tra thành viên nhanh:

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

True

Lưu ý: {} rỗng là dict, không phải set. Muốn tạo set rỗng, bạn phải viết set().

5. Nhiều vòng for và comprehension lồng nhau

Bạn có thể viết nhiều for trong một comprehension. Thứ tự đọc giống hệt khi bạn viết các vòng for lồng nhau từ trên xuống:

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')]

Comprehension trên tương đương với đoạn code sau:

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

Trường hợp phổ biến nhất là làm phẳng (flatten) một list lồng:

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]

Cách ghi nhớ: vòng for ngoài cùng đứng trước. Nếu bạn phải dừng lại đọc lần thứ hai để hiểu thứ tự, đó là dấu hiệu nên tách ra thành vòng for thường.

6. Generator expression: khi không cần cả list

Nếu thay [] bằng (), bạn sẽ có một generator expression. Python không tạo list trong bộ nhớ. Giá trị được sinh lần lượt khi có code tiêu thụ chúng.

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

1450000

Khi generator là đối số duy nhất của hàm, bạn bỏ được một cặp ngoặc. Cách này phù hợp với 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 dừng ngay khi gặp phần tử đầu tiên thỏa điều kiện. Vì vậy với dữ liệu lớn, generator vừa tiết kiệm bộ nhớ vừa nhanh hơn any([...]).

Lưu ý: Generator chỉ duyệt được một lần. Lần duyệt thứ hai sẽ không trả về gì:

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

2
0

Nếu bạn cần dùng kết quả nhiều lần, hãy dùng list comprehension.

7. Giữ cho comprehension dễ đọc

7.1. Đặt tên hàm cho logic phức tạp

Điều kiện dài làm comprehension khó đọc:

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

Bạn có thể tách điều kiện thành một hàm có tên rõ nghĩa:

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']

Lúc này comprehension đọc như một câu văn, còn logic nghiệp vụ có chỗ riêng để test.

7.2. Xuống dòng hợp lý

Comprehension không bắt buộc phải nằm trên một dòng. Trong cặp ngoặc, bạn có thể xuống dòng thoải mái:

summary = {
    order.id: order.total
    for order in orders
    if order.status != "CANCELLED"
}
print(summary)

{1: 250000, 3: 1200000, 4: 450000}

Mỗi dòng một vai trò: biểu thức, nguồn, bộ lọc.

7.3. Walrus operator để tránh tính hai lần

Đôi khi bạn vừa cần lọc vừa cần dùng kết quả của cùng một phép tính. Toán tử := (Python 3.8+) giúp bạn chỉ tính một lần:

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}

Hãy dùng walrus một cách tiết chế. Biến d sẽ “rò” ra scope bên ngoài comprehension. Ngoài ra, nhiều người đọc chưa quen cú pháp này.

8. Kiểm thử logic biến đổi dữ liệu

Comprehension thường nằm trong các hàm biến đổi dữ liệu, nên bạn nên test chúng như mọi logic khác. Trước tiên, cài pytest:

pip install pytest

Giả sử bạn có 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
    }

Và file test tương ứng:

# 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

Test đầu khẳng định hai điều: đơn bị hủy không được tính, và khách chỉ có đơn hủy không xuất hiện trong kết quả. Test thứ hai kiểm tra trường hợp biên với list rỗng.

Hàm trên đúng nhưng có độ phức tạp O(n × số khách) vì mỗi khách lại duyệt toàn bộ orders. Với dữ liệu lớn, vòng for kèm defaultdict vừa nhanh hơn vừa dễ đọc hơn:

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)

Hai test cũ vẫn pass mà không cần sửa. Đây chính là lợi ích của việc có test: bạn tự tin refactor từ comprehension sang vòng for hoặc ngược lại.

9. Khi nào KHÔNG nên dùng comprehension

• Khi cần side effect. Đừng viết [print(o) for o in orders] hay [db.save(o) for o in orders]. Cách này tạo ra một list None vô nghĩa và đánh lừa người đọc rằng bạn cần kết quả. Hãy dùng vòng for.
• Khi lồng quá hai cấp hoặc có nhiều if đan xen. Ba vòng for trong một dòng hầu như luôn khó đọc hơn phiên bản tách dòng.
• Khi cần gom nhóm hoặc tích lũy như cộng dồn hay đếm. Hãy dùng defaultdict, Counter hoặc vòng for.
• Khi cần xử lý exception cho từng phần tử. Comprehension không chứa được try/except. Bạn nên tách thành hàm riêng rồi gọi trong comprehension, hoặc dùng vòng for.
• Khi dữ liệu rất lớn mà chỉ duyệt một lần. Lúc đó hãy dùng generator expression thay cho list.

Lưu ý: Trong Python 3, biến vòng lặp của comprehension không rò ra ngoài. Sau [x for x in range(3)], tên x không tồn tại ở scope bên ngoài, khác với vòng for thường. Ngoại lệ duy nhất là biến gán bằng := như đã nói ở mục 7.3.

10. Kết luận

Comprehension giúp bạn diễn đạt ý “tạo một collection mới từ collection cũ” một cách ngắn gọn:

• [...] cho list, {k: v ...} cho dict, {...} cho set, và (...) cho generator.
• if sau for dùng để lọc, còn a if cond else b trước for dùng để biến đổi.
• Tách điều kiện phức tạp thành hàm có tên rõ nghĩa, và xuống dòng khi biểu thức dài.
• Dùng generator expression với sum, any, all để tiết kiệm bộ nhớ.
• Quay lại vòng for khi có side effect, khi cần gom nhóm, xử lý exception, hoặc khi comprehension bắt đầu khó đọc.

Quy tắc đơn giản nhất: nếu đồng nghiệp phải đọc comprehension của bạn hai lần mới hiểu, hãy viết lại bằng vòng for. Code gọn chỉ có giá trị khi vẫn dễ đọc.