Hàm Generator với yield — Tạo giá trị từng cái một để tiết kiệm bộ nhớ

Trong bài trước bạn đã thấy closure — cách niêm phong trạng thái bên trong một hàm và trả về giá trị khác mỗi lần gọi. Python có một cơ chế chuyên dụng cho loại hàm "trả về giá trị tiếp theo mỗi lần gọi" này: hàm generator. Dùng yield thay cho return, hàm tạm dừng giữa chừng để trả về một giá trị, rồi tiếp tục từ chỗ nó dừng ở lần gọi tiếp theo.

Đây là lựa chọn tuyệt vời khi bạn không muốn vật chất hóa toàn bộ một tập hợp lớn cùng lúc và muốn stream từng giá trị một — xử lý log, crunch dữ liệu lượng lớn, và các tải tương tự.

Cơ bản về yield — Trả về một giá trị mỗi lần

Một hàm có yield value trong thân chính là một hàm generator. Gọi nó như hàm thông thường không chạy thân hàm — bạn chỉ nhận về một đối tượng generator đặc biệt.

Để thực sự lấy ra một giá trị, gọi next(obj).

next() đầu tiên chạy thân hàm đến yield đầu tiên và trả về giá trị đó.

next() tiếp theo tiếp tục từ đó và chạy đến yield kế tiếp.

Khi không còn yield nào nữa, một next() khác phát sinh StopIteration.

def simple():
    yield 1
    yield 2

gen = simple()
print(type(gen))    # <class 'generator'>

print(next(gen))    # 1
print(next(gen))    # 2
print(next(gen))    # 3
# print(next(gen))  ← StopIteration (không còn yield)

Khác biệt với return

Hàm thông thường return và toàn bộ phạm vi biến mất — lần gọi tiếp theo bắt đầu từ đầu. Hàm generator thì ngược lại, tạm dừng tại yield và giữ lại biến local cùng vị trí. next() tiếp theo tiếp tục đúng chỗ bạn dừng, đó là khác biệt lớn.

Xây hàm generator yield ID đơn hàng 1, 2, 3 theo thứ tự, và kéo từng cái một bằng next().

① Định nghĩa def order_ids(): với ba dòng: yield 1 / yield 2 / yield 3.

② Tạo đối tượng generator: gen = order_ids().

③ Gọi print(next(gen)) ba lần liên tiếp và xác nhận 1 / 2 / 3 xuất hiện.

(Khi đáp án đúng, phần giải thích sẽ xuất hiện.)

Lấy giá trị bằng for — Không cần lo về StopIteration

Viết next() mỗi lần thật phiền, và bạn cũng không nên phải tự xử lý StopIteration. Thay vào đó dùng for value in generator:, và Python gọi next() ngầm bên dưới và tự thoát vòng lặp khi generator kết thúc. Đây là dạng phổ biến hơn rất nhiều.

def count_up_to(max_value):
    print("bắt đầu generator")
    for i in range(max_value):
        print(f"  trước yield: {i}")
        yield i
        print(f"  sau yield: {i}")

for v in count_up_to(3):
    print(f"nhận được: {v}")

# Luồng output:
# bắt đầu generator
#   trước yield: 0
# nhận được: 0
#   sau yield: 0
#   trước yield: 1
# ... (tiếp tục)

Dùng for để nhận dữ liệu khách hàng từ một generator stream từng cái một.

① Định nghĩa def each_customer(): và lặp for name in ["Minh", "Linh", "Hoa"]:, yield name.

② Lấy giá trị bằng for name in each_customer(): và in f"Khách tiếp theo: {name}".

Nếu cả ba xuất hiện theo thứ tự, bạn đã xong.

Khác biệt với list — Cắt giảm bộ nhớ

Khi bạn làm việc với các giá trị từ 0 đến 999_999, một list comprehension vật chất hóa cả 1.000.000 số nguyên cùng lúc trong bộ nhớ. Generator thì ngược lại, chỉ giữ giá trị hiện tại và tính cái tiếp theo theo yêu cầu. List có dung lượng vài MB; chính đối tượng generator chỉ là vài trăm byte.

import sys

MAX = 10 ** 6

# List: load tất cả vào bộ nhớ cùng lúc
data_list = [i for i in range(MAX)]
print(sys.getsizeof(data_list))
# vd. 8000056 (~8 MB)

# Generator expression: chỉ phần tử hiện tại
data_gen = (i for i in range(MAX))
print(sys.getsizeof(data_gen))
# vd. ~200 byte

# Code phía bên gọi giống nhau
for v in data_gen:
    if v > 2:
        break
    print(v)
# 0
# 1
# 2

Lối tắt Generator Expression

Đổi ngoặc vuông [ ... ] của list comprehension thành ngoặc tròn ( ... ) và bạn có generator expression. (i for i in range(1_000_000)) cho cùng hiệu ứng như hàm generator dựa trên def chỉ trong một dòng. Truyền thẳng vào sum() / max() / any() và bạn bè — chạy ngon.

Xây một generator expression (list comprehension với ngoặc tròn thay vì ngoặc vuông) cho dữ liệu giá, xác nhận kiểu bằng type(), rồi lấy giá trị từng cái bằng for.

① Tạo prices = (base * 100 for base in range(1, 6)). Mẹo là dùng ( ) thay vì [ ].

② In kiểu bằng print(type(prices)) và xác nhận <class 'generator'> xuất hiện.

③ Lấy giá trị bằng for p in prices: và in f"Giá: {p}".

Nối chuỗi Generator với yield from

Khi muốn một generator chuyển tiếp giá trị từ generator khác nguyên xi, bạn có thể viết yield from sub_generator trong một dòng thay vì for v in sub: yield v. Tiện khi muốn gộp nhiều nguồn dữ liệu vào một generator duy nhất.

Ví dụ, với một hàm stream doanh số chi nhánh Tokyo và một hàm khác cho Osaka, xếp yield from tokyo_sales() và yield from osaka_sales() cho phía bên gọi một generator trông như một dòng liên tục.

def tokyo_sales():
    yield 1200
    yield 980

def osaka_sales():
    yield 850
    yield 1340

def all_sales():
    yield from tokyo_sales()
    yield from osaka_sales()

for amount in all_sales():
    print(amount)
# 1200
# 980
# 850
# 1340

Xây một generator gộp các list tồn kho từng cửa hàng thành một dòng chảy duy nhất.

① Định nghĩa def store_a(): và for item in ["táo", "cam"]: yield item.

② Định nghĩa def store_b(): và for item in ["chuối", "nho", "dâu"]: yield item.

③ Định nghĩa def all_items(): và viết yield from store_a() rồi yield from store_b().

④ Lặp bằng for item in all_items(): print(item) để in cả 5 theo thứ tự.