[Python-ideas] "atexit" equivalent for functions in contextlib

[Nikolaus Rath]

Hello,

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?

Best,
-Nikolaus
-- 
GPG encrypted emails preferred. Key id: 0xD113FCAC3C4E599F
Fingerprint: ED31 791B 2C5C 1613 AF38 8B8A D113 FCAC 3C4E 599F

             »Time flies like an arrow, fruit flies like a Banana.«