Python "with" statement pileup

Question

I tend to use the Python "with" statement quite often. Mostly for cleaning up a directory after I have symlinked or copied some files into the directory, as the tasks still gets executed even if the python script crashes. Here is an example of on of my functions that can be used with the "with" statement.

@contextmanager
def use_symlink(orig, dest):
    os.symlink(orig, dest)
    try: 
        yield
    finally:
        os.unlink(link)

The way I use these with statements they tend to pile up quite quickly.

#Off to an adventure
with use_symlink(a, b):
    with use_symlink(c, b):
        with use_symlink(d, b):
            with working_dir(dir1):
                #do something
            with working_dir(dir2):
                #do something that creates file dir2_file1, dir2_file2
                with use_symlink(dir2_file1, b):
                   with use_symlink(dir2_file2, b):
                       with working_dir(b):
                           #Do the last thing
#Home safely

Is there a better way to do the above with the same ease and safety of the powerful "with" statement?

Answer 1

You can put multiple context managers under the same with statement:

with use_symlink(a, b), use_symlink(c, b), use_symlink(d, b):
    with working_dir(dir1):
        #do something
    with working_dir(dir2):
        #do something that creates file dir2_file1, dir2_file2
        with use_symlink(dir2_file1, b), use_symlink(dir2_file2, b), working_dir(b):
            #Do the last thing

On Python 3, you could use the contextlib.ExitStack() to manage multiple context managers:

from contextlib import ExitStack

with ExitStack() as stack:
    for combo in ((a, b), (c, b), (d, b)):
        stack.enter_context(use_symlink(*combo))
    with working_dir(dir1):
        # do something
    with working_dir(dir2):
        #do something that creates file dir2_file1, dir2_file2
        for combo in ((dir2_file1, b), (dir2_file2, b)):
            stack.enter_context(use_symlink(*combo))
        with working_dir(b):
            #Do the last thing

The use_symlink context managers are then torn down in reverse order when ExitStack() exits.