catch a break in trio async generator

Question

I have a strange error with an asynchronous trio loop. When I interrupt the loop with break during iteration, I expected the print statement "exit 2" to be printed before the statement "--- Hello ---". But when I run this example, the code after the iteration is executed before the code in the GeneratorExit exception.

import trio

class Bar:
    async def __aiter__(self):
        for x in range(10):
            try:
                yield x
            except Exception as e:
                print("ups", x)
                raise e
            except GeneratorExit as e:
                print("exit", x) 
                raise e
            else:
                print("else", x)
            finally:
                print("finally", x)
                
bar = Bar()
async for x in bar:
    if x == 2:
        break  # missing trio checkpoint?

print("--- Hello ---")

outputs:

else 0
finally 0
else 1
finally 1
--- Hello ---
exit 2
finally 2

When I put a trio.sleep(...) before the last print my use-case woks as expected. But This is not the solution I want. What can I do with my class Bar to fix this error?

bar = Bar()
async for x in bar:
    if x == 2:
        break  # missing trio checkpoint?

await trio.sleep(0.001)
print("--- Hello ---")

outputs:

else 0
finally 0
else 1
finally 1
exit 2
finally 2
--- Hello ---

Answer 1

This problem is not trio specific.

You're getting this funny result because the generator is cleaned up sometime later. "async for" is syntactic sugar for creating an iterator with __aiter__, then calling __anext__ on the result until that raises AsyncStopIteration. If that doesn't happen, GeneratorExit is thrown into the async iterator.

But how does the interpreter know that you're not going to call __anext__ any more? Answer: the iterator goes out of scope and is then garbage collected. Garbage collection doesn't happen immediately but "some time later". In your case, either when your program exits or during trio.sleep.

How to fix this:

from async_generator import aclosing ## before py3.10
bar = Bar().__aiter__()  # aiter(bar) in py3.10
async with aclosing(bar):
    async for x in bar:
        if x == 2:
            break