Composing Context Managers in Python

I freaking love python's context managers. They's excellent for managing resources that need to have a set lifetime. I find myself writing lots of code that uses context managers like tempfile.NamedTemporaryFile and tempfile.TemporaryDirectory. Recently I found myself writing code that set up a data structure on disk. The code looked similar to:

def do_some_work():
        with tempfile.TemporaryDirectory() as workdir:
            # Some more setup code
            with tempfile.NamedTemporaryFile() as data_file:
                data_file.write(b"some_data")

                # Do the actual work here with data_file

This code worked, but I wasn't happy - the function it was in was all about the 'actual work' bit. I felt that the intent of that function was getting muddied up by all the setup code. Moving the setup code to another function helps, but now we have the problem that we can't control the lifetime of those files on disk:

def setup_structure_on_disk():
        workdir = tempfile.mkdtemp()
    # Some more setup code
    tempfile = tempfile.mktemp(dir=workdir)
    tempfile.write(b"some_data")
    return tempfile

def do_some_work():
        data_file = setup_structure_on_disk()
    # Do the actual work here with data_file

    # How do we ensure those files are cleaned up?

The obviuous answer here is to turn setup_structure_on_disk into a context manager. The simplest way to do this is to use the contextlib.contextmanager decorator, like so:

@contextlib.contextmanager
def setup_structure_on_disk():
        with tempfile.TemporaryDirectory() as workdir:
            # Some more setup code
            with tempfile.NamedTemporaryFile() as tempfile:
                tempfile.write(b"some_data")
                yield tempfile

def do_some_work():
        with setup_structure_on_disk() as data_file:
        # Do the actual work here with data_file

This is quickly becoming much more readable. However, in my example I wanted to yield a contextmanager object with additional methods that manipulate the data on disk. This precludes using the contextmanager decorator without a lot of faffing about (yes, I realise you can yield whatever you want, but... ugh).

So I started writing...

class SetupDiskStructure(object):

        def __enter__(self):
            # setup / enter the context manager.
            # PROBLEM: How do I use context managers in here?
            return self

        def __exit__(self, exc_type, exc_val, exc_tb):
            # undo changes on disk.

        def some_method(self):
                """Do something useful"""


def do_some_work():
        with SetupDiskStructure() as manager:
                manager.some_method()

The problem I face here is that I'd really like to be able to tell python "I'm using this context manager, please call it's __exit__ when you call mine". I couldn't think of a nice way to do that. I wanted to composse my SetupDiskStructure context manager of serveral other context managers.

Enter ExitStack

I have used contextlib.ExitStack several times before. It's extremely useful any time you have a variable number of context managers that you want to control in one block. For example, if you want to write a string to every file in a particular directory, you can do something like this:

file_list = os.listdir('/path/to/directory')
with contextlib.ExitStack() as context_stack:
        open_files = [context_stack.enter_context(open(f)) for f in file_list]

We now have a list of file object that will be closed when we leave the 'with' block.

I had never considered (ab)using ExitStack to help me solve my problem above. In the end I ended up writing code that looked similar to this:

class SetupDiskStructure(ExitStack):

        def __enter__(self):
                super().__enter__()
                self.workdir = self.enter_context(tempfile.TemporaryDirectory())
                self.data_file = self.enter_context(tempfile.NamedTemporaryFile(dir=self.workdir))
                self.data_file.write(b'Hello')
            return self

        def some_method(self):
                """Do something useful"""

def do_some_work():
        with SetupDiskStructure() as manager:
                # can still call methods on the context manager:
                manager.some_method()
                # plus workdir and data_file members:
                print("Working directory is: %s" % manager.workdir)
                print("data_file path is: %s" % manager.data_file.name)

I found it... pleasing that I can delegate 100% of the cleanups to the context managers I'm calling - no need to write my own __exit__ at all (of course, I could, if I wanted to).