Python生成器中避免yield from后的额外next调用方法

问题根源 #

当yield from DUMP_DATA执行完毕后，原代码通过continue进入下一次循环，此时会执行in_value = yield out_value产出None，但这个产出的None未被消费，生成器就暂停在该位置。因此必须调用一次next(g)来消费这个None，生成器才能回到可接收send的状态。

解决方案 #

调整异常分支的逻辑，在yield from结束后，手动完成一次yield并接收后续send的值，直接将生成器推进到可接收状态，无需额外的next调用。

修改后的生成器代码

from collections.abc import Generator

DUMP_DATA = 5, 6, 7

class DumpData(Exception):
    """Exception used to indicate to yield from DUMP_DATA."""

def sample_gen() -> Generator[int | None, int, None]:
    out_value: int | None = None
    while True:
        try:
            in_value = yield out_value
        except DumpData:
            yield len(DUMP_DATA)
            yield from DUMP_DATA
            out_value = None
            # 手动执行yield并接收send值，直接进入可接收状态
            in_value = yield out_value
            out_value = in_value
            continue
        out_value = in_value

对应的调用代码（无需额外next）

def main() -> None:
    g = sample_gen()
    next(g)  # 初始化生成器
    assert g.send(1) == 1
    assert g.send(2) == 2

    # 导出数据
    num_data = g.throw(DumpData)
    data = tuple(next(g) for _ in range(num_data))
    assert data == DUMP_DATA

    # 无需额外next调用，直接send即可
    assert g.send(3) == 3

流程说明 #

1. 触发DumpData异常后，先产出数据长度，再通过yield from产出所有数据；
2. 设置out_value = None后，执行in_value = yield out_value，产出None并暂停；
3. 调用g.send(3)时，3会被赋值给in_value，out_value随之更新为3；
4. 通过continue进入下一次循环，执行in_value = yield out_value产出3，这正是send(3)的返回值，符合断言要求。

内容的提问来源于stack exchange，提问作者Intrastellar Explorer