Register function to be executed on function exit

Question

Python has the atexit built-in API for registering functions to be executed on program exit.

In my case, I have an API function that can be imported and used, which creates a temporary directory that needs to be cleaned-up when the function returns (or raises an exception).

Rather than using try-finally, is there any conventional way to register a function to be executed on function exit?

Answer 1

The simplest way to do this reusably is just to use a decorator:

def cleanup(action):
    def decorator(func):
        def wrapper(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            finally:
                action()
        return wrapper
    return decorator

For example:

>>> def postrun():
...     print("afterwards")
... 
>>> @cleanup(postrun) 
... def divider(x, y):
...     return x / y
... 
>>> divider(2, 3)
afterwards
0.6666666666666666
>>> divider(2, 0)
afterwards
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/jsharpe/workspace/perigon/test.py", line 5, in wrapper
    return func(*args, **kwargs)
  File "<stdin>", line 3, in divider
ZeroDivisionError: division by zero

You could also pass args and kwargs to action if you think they're going to be needed for the clean-up.