2.31. Iterator và generator¶

Vòng lặp for đã làm nhiều việc hơn vẻ ngoài của nó. Trang này đề cập đến giao thức iterator mà nó hoạt động dựa trên, và từ khóa yield cho phép bạn tự xây dựng iterator.

2.31.1. Giao thức iterator¶

Mỗi đối tượng có thể được lặp qua đều triển khai hai phương thức:

__iter__() -- trả về một iterator qua các phần tử của đối tượng.
__next__() -- trên iterator, trả về phần tử tiếp theo hoặc raise StopIteration khi không còn phần tử nào.

Hàm built-in iter() gọi __iter__; next() gọi __next__. Duyệt qua một danh sách thủ công:

it = iter([10, 20, 30])
print(next(it))    # 10
print(next(it))    # 20
print(next(it))    # 30
print(next(it))    # raises StopIteration

A for loop calls __iter__ once to get an iterator, then calls __next__ repeatedly until StopIteration ends the loop. — `for` là cách viết tắt của "gọi `__iter__` một lần, sau đó lặp trên `__next__` cho đến `StopIteration`".¶

Những gì for x in items: thực sự làm:

_it = iter(items)
while True:
    try:
        x = next(_it)
    except StopIteration:
        break
    # ... loop body ...

Mỗi list, tuple, string, dict, set, file object và generator đều đã triển khai __iter__ và __next__ -- đó là lý do tại sao tất cả chúng đều hoạt động với for.

2.31.2. yield và hàm generator¶

Một hàm có chứa câu lệnh yield là một hàm generator. Gọi nó sẽ không chạy phần thân; nó trả về một đối tượng generator (một iterator) chạy phần thân mỗi lần một yield:

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

for value in count_up_to(3):
    print(value)

Kết quả:

0
1
2

Mỗi lần gọi next() tiếp tục hàm cho đến yield tiếp theo, chuyển giá trị đó cho caller và tạm dừng ở đó. Trạng thái cục bộ (i trong trường hợp này) được giữ nguyên giữa các lần tiếp tục.

Sequence diagram with two lifelines (caller and generator body). The caller calls count_up_to(3), which creates a generator without running the body. Each subsequent next() runs the body until the next yield, returns the yielded value, and pauses. The fourth next() falls off the end and raises StopIteration. The variable i is preserved across pauses. — `next()` chạy phần thân cho đến `yield` tiếp theo, chuyển giá trị lại và tạm dừng. Trạng thái cục bộ tồn tại qua lần tạm dừng.¶

Generator là cách dễ nhất để tạo ra một chuỗi một cách lười biếng -- không có danh sách nào được xây dựng, các phần tử chỉ được tính toán khi consumer yêu cầu, và hàm có thể yield các phần tử mãi mãi nếu muốn.

2.31.3. Pipeline lười biếng¶

Generator kết hợp tốt với nhau. Đầu ra của một generator có thể cấp cho generator khác:

def numbers():
    i = 0
    while True:
        yield i
        i += 1

def squares(source):
    for x in source:
        yield x * x

pipeline = squares(numbers())

for v in pipeline:
    if v > 100:
        break
    print(v)

Các giá trị chảy qua pipeline từng cái một -- không có danh sách trung gian, không có giới hạn trên được tích hợp vào numbers, và consumer (for v in pipeline) quyết định khi nào dừng.

Three boxes left-to-right: numbers(), squares(source), and the for-v-in-pipeline consumer. Three cycles are drawn beneath. In each cycle, the consumer sends a pull request leftward to squares, which sends a pull leftward to numbers; numbers yields a value rightward to squares, which yields its squared value rightward to the consumer. — Mỗi `next()` trên consumer kích hoạt một lần kéo qua chuỗi; các giá trị chỉ tồn tại khi có thứ gì đó yêu cầu chúng.¶

2.31.3.1. `yield from`¶

Một vòng lặp kéo các phần tử từ một iterable khác và yield từng phần tử là đủ phổ biến để Python cung cấp một phím tắt. Biểu thức yield from iter yield từng giá trị mà iterable tạo ra, theo thứ tự -- như thể generator có vòng lặp for x in iter: yield x nội tuyến:

def chain(*sources):
    for source in sources:
        yield from source

for v in chain([1, 2, 3], (4, 5), "abc"):
    print(v)

Kết quả:

yield from hoàn toàn tương đương với vòng lặp for tường minh, chỉ ngắn hơn, và nó truyền StopIteration từ iterable bên trong lên generator bên ngoài một cách sạch sẽ -- hữu ích khi nối nhiều generator đầu-đuôi.

2.31.3.2. Khi `yield` cạn kiệt¶

Kết thúc phần thân hàm generator (hoặc đạt đến return tường minh) sẽ tự động raise StopIteration. Không cần raise thủ công; vòng lặp for xung quanh sẽ nhận thấy và kết thúc.

Hãy dùng generator khi code sản xuất được viết tự nhiên như một vòng lặp với một vài điểm yield; hãy dùng list comprehension đơn giản khi bạn thực sự cần toàn bộ chuỗi trong bộ nhớ.