使用break时，在python生成器中的yield语句之后运行代码

Question

我有一个简单的生成器：

def counter(n): 
counter = 0
while counter <= n:
counter += 1
yield counter
print(f"Nice! You counted to {counter}")
def test_counter(): 
for count in counter(6): 
print(f"count={count}")
if count > 3: 
break

在本例中，一旦计数达到3，就会出现中断。然而，代码的唯一问题是break终止了yield语句处的生成器执行。它不会在yield语句之后运行print语句。有没有一种方法可以在发生中断时在yield语句之后的生成器中执行代码？

示例运行：

# What happens: 
>>> test_counter()
count=1
count=2
count=3
count=4
# What I'd like to see: 
>>> test_counter()
count=1
count=2
count=3
count=4
Nice! You counted to 4

Answer 1

您必须将print放在finally块中：

def counter(n): 
try:
counter = 0
while counter <= n:
counter += 1
yield counter
finally:
print(f"Nice! You counted to {counter}")
def test_counter(): 
for count in counter(6): 
print(f"count={count}")
if count > 3: 
break

Answer 2

您可以使用生成器的send方法：

def counter(n): 
counter = 0
while counter <= n:
counter += 1
should_continue = yield counter
if should_continue == False:
break
print(f"Nice! You counted to {counter}")
def test_counter(): 
gen = counter(6)
for count in gen: 
print(f"count={count}")
if count > 3: 
try:
gen.send(False)
except StopIteration:
break

现在test_counter()按预期工作。