如何优雅的获取生成器的 return 值？

Question

from loguru import logger


def func():
    yield 1
    return 2


for i in func():
    logger.debug(i)

使用 for 循环遍历的方式好像无法获取 return 值？

from loguru import logger


def func():
    yield 1
    return 2


x = func()

r = x.send(None)
logger.debug(r)
r = x.send(100)
logger.debug(r)

手工 send 的方式，获取的 return 值异常里面

2022-06-23 09:56:53.892 | DEBUG    | __main__:<module>:12 - 1
Traceback (most recent call last):
  File "/Users/ponponon/Desktop/code/me/ideaboom/010.py", line 13, in <module>
    r = x.send(100)
StopIteration: 2

要获取这个返回值，还要写一个 try except 结构，太丑陋了

有什么优雅的解决方案吗？

Answer 1

生成器 return 的值就是跟 StopIteration 一起抛出来的，for 通过捕获这个异常来结束循环

for x in f():
    print(x)
# 相当于
x_gen = f()
while True:
    try:
        x = next(x_gen)
        print(x)
    except StopIteration:
        break

通过捕获异常获取返回值，可以写个方法

def get_return_value(f):
    gen = f()
    while True:
        try:
            next(gen)
        except StopIteration as e:
            return e.value


def f():
    yield 1
    return 2


print(get_return_value(f))  # 2