diff --git a/Lib/contextlib.py b/Lib/contextlib.py --- a/Lib/contextlib.py +++ b/Lib/contextlib.py @@ -138,3 +138,119 @@ return self.thing def __exit__(self, *exc_info): self.thing.close() + + +class CleanupManager(object): + '''Context manager to automatically call registered "cleanup" + functions at the end of a block. + + This context manager is typically used to manage resources that + need to be explicitly freed but don't come with their own context + manager. It also avoids deeply nested with try..finally blocks if + several resources need to be managed. + + Example 1: + + with CleanupManager() as mngr: + tmpdir = tempfile.mkdtemp() + mngr.register(shutil.rmtree, args=(tmpdir,)) + + + is equivalent to + + try: + tmpdir = tempfile.mkdtemp() + + finally: + shutil.rmtree(tmpdir) + + Example 2: + + allocate_res1() + try: + # do stuff + allocate_res2() + try: + # do stuff + allocate_res3() + try: + # do stuff + finally: + cleanup_res3() + finally: + cleanup_res2() + finally: + cleanup_res1() + + can also be written as: + + with CleanupManager() as mngr: + allocate_res1() + mngr.register(cleanup_res1) + # do stuff + allocate_res2() + mngr.register(cleanup_res2) + # do stuff + allocate_res3() + mngr.register(cleanup_res3) + # do stuff + ''' + + def __init__(self): + self.cleanup_callbacks = list() + + def register(self, callback, *args, **kwargs): + '''Register *callback* as a cleanup function + + This method instructs the `CleanupManager` instance to add + *callback* to the list of functions to be called when control + leaves the managed block. The function will be called with + arguments *args* and keywords arguments *kwargs*. + ''' + + self.cleanup_callbacks.append((callback, args, kwargs)) + + def unregister(self, callback, *args, **kwargs): + '''Remove *callback* from the list of cleanup functions + + This method instructs the `CleanupManager` instance to remove + *callback* from the list of functions to be called when control + leaves the managed block. + + This method removes all calls to *callback* with the specified + arguments. To remove all calls to *callback* disregarding the + arguments, use the `unregister_all` functions. + ''' + + while (callback, args, kwargs) in self._cleanup_callbacks: + self._cleanup_callbacks.remove((callback, args, kwargs)) + + def unregister_all(self, callback): + '''Remove *callback* from the list of cleanup functions + + This method instructs the `CleanupManager` instance to remove + *callback* from the list of functions to be called when control + leaves the managed block. + + This method removes all calls to *callback*, disregarding the + arguments passed to `register`. + ''' + + for i in [ i for (i, el) in self.cleanup_callbacks + if el[0] == callback ]: + del self.cleanup_callbacks[i] + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + self._next_callback() + + def _next_callback(self): + if self.cleanup_callbacks: + (callback, args, kwargs) = self.cleanup_callbacks.pop() + try: + callback(*args, **kwargs) + finally: + self._next_callback() +