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

run code after yield statement in python generator when using break

我有一个简单的生成器:

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

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

您必须将 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