The other day, I was looking for an "atexit" equivalent at the function level. I was hoping to replace code like this:
def my_function(bla, fasel): with ExitStack() as cm: ... cm.push(...) ...
with something like:
@handle_on_return def my_function(bla, fasel, on_return): ... on_return.push(...) ...
It seems that contextlib *almost* offers a way to do this with ExitStack and ContextDecorator - but as far as I can tell the final piece is missing, because ContextDecorator does not offer a way to pass the context manager to the decorated function. However, the following decorator does the job:
def handle_on_return(fn): @functools.wraps(fn) def wrapper(*a, **kw): with contextlib.ExitStack() as on_return: kw['on_return'] = on_return return fn(*a, **kw) return wrapper
It's not a lot of code, but my feeling is that not anyone who might be interested in an "on_return" functionality would be able to come up with this right away.
Would it make sense to add something along these lines to contextlib? Maybe instead of a new decorator, ContextDecorator could also take an additional keyword argument that tells it to pass the context manager to the decorated function?