使用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()按预期工作