List comprehension và generator: viết code Python gọn hơn

List comprehension là một trong những đặc trưng dễ nhận ra nhất của code Python. Dùng đúng chỗ, nó làm code ngắn và rõ. Dùng quá tay, nó biến code thành thứ không ai đọc nổi.

Cú pháp cơ bản

# Cách viết thông thường
binh_phuong = []
for x in range(10):
    binh_phuong.append(x ** 2)

# List comprehension
binh_phuong = [x ** 2 for x in range(10)]

Thêm điều kiện lọc:

chan = [x for x in range(20) if x % 2 == 0]

# Điều kiện ở vị trí giá trị (dùng if/else)
nhan = ["chẵn" if x % 2 == 0 else "lẻ" for x in range(5)]

Chú ý vị trí khác nhau: if đứng cuối là để lọc, còn if/else đứng trước for là để chọn giá trị.

Comprehension cho dict và set

# Dict
gia = {"táo": 20000, "cam": 35000, "xoài": 50000}
gia_sau_giam = {ten: g * 0.9 for ten, g in gia.items()}

# Đảo khóa và giá trị
dao = {v: k for k, v in gia.items()}

# Set — tự loại trùng
do_dai = {len(tu) for tu in ["python", "java", "go", "rust"]}
# {2, 4, 6}

Khi nào KHÔNG nên dùng

Comprehension lồng nhiều tầng nhanh chóng trở nên khó đọc:

# Khó hiểu
ket_qua = [xu_ly(x) for sub in du_lieu for x in sub if kiem_tra(x) and x.hop_le]

# Rõ ràng hơn
ket_qua = []
for sub in du_lieu:
    for x in sub:
        if kiem_tra(x) and x.hop_le:
            ket_qua.append(xu_ly(x))

Quy tắc thực dụng: nếu comprehension không vừa một dòng dễ đọc, hoặc có quá hai mệnh đề for, hãy viết vòng lặp thường.

Cũng đừng dùng comprehension chỉ để gây tác dụng phụ:

[print(x) for x in danh_sach]     # sai mục đích — tạo list rác rồi vứt đi

for x in danh_sach:               # đúng
    print(x)

Generator — xử lý dữ liệu lớn

Comprehension tạo ra toàn bộ danh sách trong bộ nhớ. Với dữ liệu lớn, đây là vấn đề:

tong = sum([x ** 2 for x in range(10_000_000)])    # tạo list 10 triệu phần tử

Chỉ cần đổi ngoặc vuông thành ngoặc tròn, bạn có generator — sinh từng giá trị khi cần, không giữ tất cả trong bộ nhớ:

tong = sum(x ** 2 for x in range(10_000_000))      # gần như không tốn bộ nhớ

Viết generator bằng yield

def doc_file_lon(duong_dan):
    with open(duong_dan, encoding="utf-8") as f:
        for dong in f:
            yield dong.strip()

for dong in doc_file_lon("nhat_ky.log"):
    if "ERROR" in dong:
        print(dong)

Hàm này đọc được file hàng chục gigabyte trên máy chỉ có vài gigabyte RAM, vì tại mỗi thời điểm nó chỉ giữ một dòng.

Ghép nhiều generator thành chuỗi xử lý

def doc_dong(duong_dan):
    with open(duong_dan, encoding="utf-8") as f:
        yield from f

def bo_dong_trong(dong_iter):
    for dong in dong_iter:
        if dong.strip():
            yield dong

def chi_lay_loi(dong_iter):
    for dong in dong_iter:
        if "ERROR" in dong:
            yield dong

def phan_tich(dong_iter):
    for dong in dong_iter:
        phan = dong.split(" | ")
        yield {"thoi_gian": phan[0], "thong_bao": phan[-1]}

# Ghép lại thành một chuỗi xử lý
duong_ong = phan_tich(chi_lay_loi(bo_dong_trong(doc_dong("app.log"))))

for ban_ghi in duong_ong:
    print(ban_ghi["thoi_gian"], ban_ghi["thong_bao"])

Không có bước nào trong chuỗi này nạp toàn bộ file vào bộ nhớ. Dữ liệu chảy qua từng dòng một, và mỗi hàm chỉ làm đúng một việc — dễ đọc, dễ kiểm thử riêng lẻ.

Generator chỉ duyệt được một lần

Đây là điểm khác biệt quan trọng so với list:

gen = (x for x in range(5))
print(list(gen))     # [0, 1, 2, 3, 4]
print(list(gen))     # [] — đã cạn

Nếu cần duyệt nhiều lần, hãy chuyển sang list, hoặc viết một hàm trả về generator mới mỗi lần gọi.

itertools

Thư viện chuẩn có sẵn nhiều công cụ làm việc với generator:

from itertools import islice, chain, groupby, batched

# Lấy 10 phần tử đầu từ một nguồn vô hạn
dau_tien = list(islice(nguon_du_lieu, 10))

# Nối nhiều iterable
tat_ca = chain(danh_sach_1, danh_sach_2, danh_sach_3)

# Chia thành từng lô (Python 3.12+)
for lo in batched(range(10), 3):
    print(lo)     # (0,1,2), (3,4,5), (6,7,8), (9,)

# Gom nhóm — lưu ý phải sắp xếp trước
du_lieu = sorted(nhan_vien, key=lambda x: x["phong_ban"])
for phong, nhom in groupby(du_lieu, key=lambda x: x["phong_ban"]):
    print(phong, len(list(nhom)))

groupby chỉ gom các phần tử liền kề có cùng khóa. Quên sắp xếp trước là lỗi rất hay gặp và kết quả sẽ sai một cách âm thầm.

Chọn giữa list và generator

Dùng list khi dữ liệu nhỏ, cần truy cập ngẫu nhiên, cần biết độ dài, hoặc phải duyệt nhiều lần.
Dùng generator khi dữ liệu lớn hoặc không biết trước độ dài, chỉ duyệt một lần, hoặc muốn xử lý theo luồng.

Một chỉ dấu đơn giản: nếu bạn viết comprehension rồi truyền thẳng vào sum, any, all, max hay một vòng for, hãy bỏ cặp ngoặc vuông đi. Bạn không cần cái list đó.