使用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 时有一个中断。但是代码的唯一问题是中断在 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-block 中：

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