Alexander Kondratskiy
Alexander Kondratskiy

Reputation: 4275

Run code on coroutine close()

I am writing code that uses coroutines heavily, and I want reliable behavior on shutdown.

Say I have a coroutine and a context manager:

from contextlib import contextmanager

@contextmanager
def print_context_manager(text):
    print("Enter", text)
    yield
    print("Exit", text)

def coro():
    with print_context_manager("coro"):
        while True:
            print("Loop", (yield))

I could use it like this:

c = coro()
next(c)
c.send("Hello ")
c.send("World!")
c.close()

Unfortunately, as far as I can tell, there is no way to execute my own code on c.close(). In particular the context manager in coroutine never prints "Exit coro"

What's the point of context managers in coroutines? Do I have to manually come up with a way to signal the end of a stream? What's the point of close() then?

See this example: https://repl.it/M0XI/0

Upvotes: 3

Views: 2390

Answers (1)

user2357112
user2357112

Reputation: 280973

Your context manager has a bug. Correct it, and it will automatically perform cleanup when the coroutine is closed.

Closing a coroutine works by raising GeneratorExit at the point where the coroutine is suspended. If the code in the with raises an exception, @contextlib.contextmanager raises that exception at the point of the yield. Your context manager doesn't deal with that, so the exception prevents the cleanup from running.

You need to wrap the yield in a try-finally and do cleanup in the finally if you want the cleanup to run even on an exception:

@contextmanager
def print_context_manager(text):
    print("Enter", text)
    try:
        yield
    finally:
        print("Exit", text)

Upvotes: 2

Related Questions