Generator trong Python: yield, lặp lười và generator expression

Một hàm biết tạm dừng

Generator trông giống một hàm thông thường, nhưng thay vì tính toàn bộ kết quả rồi trả về, nó yield từng giá trị một, tạm dừng giữa các lần yield cho đến khi nơi đang hỏi muốn lấy giá trị tiếp theo.

Ví dụ đơn giản nhất:

def count_up_to(n):
    current = 1
    while current <= n:
        yield current
        current += 1

for number in count_up_to(5):
    print(number)

Để ý yield thay vì return. Lần đầu tiên for hỏi xin một giá trị, Python chạy phần thân hàm cho đến khi gặp yield 1. Hàm tạm dừng ngay tại đó, trao 1 lại cho vòng lặp và nhớ chính xác chỗ nó đã dừng, kể cả các biến. Lần lặp tiếp theo tiếp tục từ chỗ cũ: current += 1, quay lại while, yield 2. Cứ thế cho đến khi điều kiện của vòng lặp sai, lúc đó generator đơn giản là dừng lại.

Tạm dừng rồi tiếp tục: toàn bộ bí quyết nằm ở đó.

Sao không tạo luôn một list?

Vì phiên bản dùng list phải cấp phát tất cả giá trị ngay từ đầu:

def count_up_to_list(n):
    return list(range(1, n + 1))

print(count_up_to_list(5))

Với 5 phần tử thì ổn. Giờ hãy tưởng tượng bạn muốn 50 triệu số nguyên, và bạn chỉ quan tâm đến số đầu tiên thỏa mãn một điều kiện nào đó. Phiên bản list cấp phát 50 triệu số rồi bạn vứt đi gần hết. Phiên bản generator chỉ tạo đúng số lượng mà nơi dùng cần. Khi vòng lặp for tìm được thứ nó muốn và break, generator đơn giản là dừng lại.

Đó là điều đáng ghi nhớ: generator cho phép bạn viết code lặp mà không cần quyết định trước mình sẽ cần bao nhiêu kết quả.

Generator expression

Nếu bạn đã từng viết list comprehension, bạn đã biết cú pháp rồi: chỉ cần đổi ngoặc vuông thành ngoặc tròn:

# List comprehension: allocates the whole result.
squares_list = [n * n for n in range(10)]

# Generator expression: computes on demand.
squares_gen = (n * n for n in range(10))

print(squares_list)
print(squares_gen)

for s in squares_gen:
    print(s)

squares_gen chưa tính gì cả. Nó chỉ là một công thức. Lặp qua nó sẽ chạy công thức từng bước một.

Generator expression rất hợp để làm đối số cho các hàm nhận một iterable:

numbers = [3, 8, 1, 5, 4, 9, 2]

print(sum(n * n for n in numbers))
print(max(n for n in numbers if n % 2 == 0))
print(any(n > 7 for n in numbers))

Không có list trung gian nào. sum, max và any đọc từng giá trị một, đúng như cách chúng muốn.

Đọc một tệp lớn theo từng dòng

Đây là trường hợp thực tế kinh điển của generator: xử lý một tệp quá lớn để nạp vào bộ nhớ:

def parse_log_lines(path):
    with open(path) as f:
        for line in f:
            if line.startswith("ERROR"):
                yield line.rstrip()

for error in parse_log_lines("app.log"):
    print(error)

Tệp được đọc theo kiểu lười. Mỗi lần gọi generator sẽ lấy một dòng từ ổ đĩa, lọc nó rồi yield. Lượng bộ nhớ dùng giữ nguyên bất kể kích thước tệp.

Dùng một lần là hết

Generator chỉ duyệt được một lượt. Sau khi bạn lặp đến cuối, nó bị cạn:

def first_three():
    yield 1
    yield 2
    yield 3

gen = first_three()

for n in gen:
    print("first time:", n)

for n in gen:
    print("second time:", n)

Vòng lặp thứ hai không in ra gì. Generator không còn gì để trả.

Nếu cần lặp nhiều lần, hoặc bạn gọi lại hàm generator để có một generator mới, hoặc chuyển dãy thành list bằng list(...) rồi lặp qua list đó bao nhiêu lần tùy ý. Chọn dựa trên chi phí: tạo lại cũng ổn nếu việc tính toán nhẹ; dùng list cũng ổn nếu dãy nhỏ.

next() và lặp thủ công

Bạn không bắt buộc phải dùng vòng lặp for. next() lấy ra từng giá trị một:

def ticks():
    yield "first"
    yield "second"
    yield "third"

gen = ticks()

print(next(gen))
print(next(gen))
print(next(gen))

# A fourth call would raise StopIteration.

StopIteration là cách generator báo hiệu "Tôi xong rồi." Vòng lặp for bắt nó một cách âm thầm. Khi tự viết code thủ công, bạn có thể truyền giá trị mặc định cho next(gen, default) để tránh ngoại lệ này.

Generator vô hạn

Vì giá trị được tạo khi cần, một generator có thể biểu diễn một dãy không có điểm kết thúc, miễn là nơi dùng biết lúc nào thôi hỏi:

def naturals():
    n = 1
    while True:
        yield n
        n += 1

for n in naturals():
    if n > 5:
        break
    print(n)

while True với một yield bên trong không làm treo chương trình. Nó chỉ có nghĩa là "nếu còn ai hỏi thì cứ tiếp tục tạo." Nơi dùng quyết định khi nào dừng.

Mẫu này xuất hiện trong xử lý dữ liệu dạng luồng, vòng lặp sự kiện và bất cứ đâu bạn lấy giá trị từ một nguồn không có độ dài xác định.

yield from: ủy quyền cho một iterable khác

Nếu generator của bạn muốn yield mọi giá trị từ một iterable khác, yield from làm việc đó trong một dòng:

def chain(*iterables):
    for it in iterables:
        yield from it

combined = chain([1, 2, 3], (4, 5, 6), "ab")
for x in combined:
    print(x)

Không có yield from, bạn sẽ phải viết một vòng lặp for lồng nhau với yield x bên trong. Nó cũng chuyển tiếp đúng các lời gọi send() và throw() nếu bạn có dùng đến, nhưng với code thường ngày, cứ hiểu nó là "yield mọi giá trị từ thứ này."

Khi nào nên dùng generator

Ba dấu hiệu cho thấy generator là công cụ phù hợp:

1. Dãy giá trị lớn, có thể vô hạn, hoặc tốn kém nếu tạo ra toàn bộ.
2. Nơi dùng có thể dừng trước khi đến cuối (ví dụ break ngay khi tìm thấy kết quả đầu tiên).
3. Bạn muốn nối chuỗi các phép biến đổi như lọc, ánh xạ, lấy vài phần tử, mà không tạo các list trung gian.

Và khi không nên dùng:

• Bạn cần truy cập ngẫu nhiên (seq[42]). Generator chỉ đi về phía trước.
• Bạn cần lặp cùng một dãy nhiều lần. Hãy dùng list.
• Dãy nhỏ và bạn đã có sẵn nó. List comprehension đơn giản hơn.

Generator, list comprehension và list thông thường đều là câu trả lời đúng cho những việc khác nhau. Kỹ năng ở đây là chọn được cái phù hợp mà không cần nghĩ quá lâu, và cách nhanh nhất để có được bản năng đó là để ý, với mỗi vòng lặp bạn viết, xem "tạo hết mọi thứ trước" hay "tạo từng cái một" hợp hơn.

Tiếp theo: context manager chuyên sâu

Bạn đã thấy hầu hết các cách viết Python dùng cho việc lặp. Tiếp theo là context manager, tức câu lệnh with, và chúng kết hợp rất tốt với generator khi truyền dữ liệu dạng luồng từ tệp và kết nối mạng.